Play with RSL-RL agent.
(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg)
| 83 | |
| 84 | @hydra_task_config(args_cli.task, args_cli.agent) |
| 85 | def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg): |
| 86 | """Play with RSL-RL agent.""" |
| 87 | # grab task name for checkpoint path |
| 88 | task_name = args_cli.task.split(":")[-1] |
| 89 | train_task_name = task_name.replace("-Play", "") |
| 90 | |
| 91 | # override configurations with non-hydra CLI arguments |
| 92 | agent_cfg: RslRlBaseRunnerCfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli) |
| 93 | env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs |
| 94 | |
| 95 | # set the environment seed |
| 96 | # note: certain randomizations occur in the environment initialization so we set the seed here |
| 97 | env_cfg.seed = agent_cfg.seed |
| 98 | env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device |
| 99 | |
| 100 | # specify directory for logging experiments |
| 101 | log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name) |
| 102 | log_root_path = os.path.abspath(log_root_path) |
| 103 | print(f"[INFO] Loading experiment from directory: {log_root_path}") |
| 104 | if args_cli.use_pretrained_checkpoint: |
| 105 | resume_path = get_published_pretrained_checkpoint("rsl_rl", train_task_name) |
| 106 | if not resume_path: |
| 107 | print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.") |
| 108 | return |
| 109 | elif args_cli.checkpoint: |
| 110 | resume_path = retrieve_file_path(args_cli.checkpoint) |
| 111 | else: |
| 112 | resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) |
| 113 | |
| 114 | log_dir = os.path.dirname(resume_path) |
| 115 | |
| 116 | # set the log directory for the environment (works for all environment types) |
| 117 | env_cfg.log_dir = log_dir |
| 118 | |
| 119 | # create isaac environment |
| 120 | env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) |
| 121 | |
| 122 | # convert to single-agent instance if required by the RL algorithm |
| 123 | if isinstance(env.unwrapped, DirectMARLEnv): |
| 124 | env = multi_agent_to_single_agent(env) |
| 125 | |
| 126 | # wrap for video recording |
| 127 | if args_cli.video: |
| 128 | video_kwargs = { |
| 129 | "video_folder": os.path.join(log_dir, "videos", "play"), |
| 130 | "step_trigger": lambda step: step == 0, |
| 131 | "video_length": args_cli.video_length, |
| 132 | "disable_logger": True, |
| 133 | } |
| 134 | print("[INFO] Recording videos during training.") |
| 135 | print_dict(video_kwargs, nesting=4) |
| 136 | env = gym.wrappers.RecordVideo(env, **video_kwargs) |
| 137 | |
| 138 | # wrap around environment for rsl-rl |
| 139 | env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) |
| 140 | |
| 141 | print(f"[INFO]: Loading model checkpoint from: {resume_path}") |
| 142 | # load previously trained model |