|
|
from navsim.agents.abstract_agent import AbstractAgent |
|
|
from omegaconf import DictConfig |
|
|
import hydra |
|
|
from hydra.utils import instantiate |
|
|
import argparse |
|
|
import os |
|
|
from hugsim.visualize import to_video, save_frame |
|
|
from hugsim.dataparser import parse_raw |
|
|
import torch |
|
|
from hugsim_client import HugsimClient |
|
|
|
|
|
CONFIG_PATH = "navsim/planning/script/config/HUGSIM" |
|
|
CONFIG_NAME = "transfuser" |
|
|
|
|
|
hugsim_client = HugsimClient() |
|
|
hugsim_client.reset_env() |
|
|
|
|
|
def get_opts(): |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument('--output', type=str, required=True) |
|
|
return parser.parse_args() |
|
|
|
|
|
|
|
|
@hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME, version_base=None) |
|
|
def main(cfg: DictConfig) -> None: |
|
|
cfg.agent.config.latent = True |
|
|
cfg.agent.checkpoint_path = "./ckpts/ltf_seed_0.ckpt" |
|
|
print(cfg) |
|
|
agent: AbstractAgent = instantiate(cfg.agent) |
|
|
agent.initialize() |
|
|
agent = agent.cuda() |
|
|
|
|
|
print('Ready for recieving') |
|
|
cnt = 0 |
|
|
output_folder = os.path.join(cfg.output, 'ltf') |
|
|
os.makedirs(output_folder, exist_ok=True) |
|
|
|
|
|
env_state = hugsim_client.get_current_state() |
|
|
while True: |
|
|
if env_state.done: |
|
|
return |
|
|
|
|
|
data = parse_raw((env_state.state.obs, env_state.state.info)) |
|
|
del data['raw_imgs'] |
|
|
|
|
|
try: |
|
|
with torch.no_grad(): |
|
|
traj = agent.compute_trajectory(data['input']) |
|
|
except RuntimeError as e: |
|
|
traj = None |
|
|
print(e) |
|
|
|
|
|
if traj is None: |
|
|
print('Waiting for visualize tasks...') |
|
|
return |
|
|
|
|
|
|
|
|
imu_way_points = traj.poses[:, :3] |
|
|
way_points = imu_way_points[:, [1, 0, 2]] |
|
|
way_points[:, 0] *= -1 |
|
|
|
|
|
env_state = hugsim_client.execute_action(way_points[:, :2]) |
|
|
print('sent') |
|
|
cnt += 1 |
|
|
if env_state.cur_scene_done: |
|
|
print('Scene done') |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
|
|
|
main() |
