(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: dict)
| 125 | |
| 126 | @hydra_task_config(args_cli.task, args_cli.agent) |
| 127 | def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: dict): |
| 128 | task_name = args_cli.task.split(":")[-1] |
| 129 | train_task_name = task_name.replace("-Play", "") |
| 130 | |
| 131 | env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs |
| 132 | env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device |
| 133 | |
| 134 | if args_cli.seed == -1: |
| 135 | args_cli.seed = random.randint(0, 10000) |
| 136 | agent_cfg["params"]["seed"] = args_cli.seed if args_cli.seed is not None else agent_cfg["params"]["seed"] |
| 137 | env_cfg.seed = agent_cfg["params"]["seed"] |
| 138 | |
| 139 | log_root_path = os.path.abspath(os.path.join("logs", "rl_games", agent_cfg["params"]["config"]["name"])) |
| 140 | print(f"[INFO] Loading experiment from directory: {log_root_path}") |
| 141 | if args_cli.use_pretrained_checkpoint: |
| 142 | resume_path = get_published_pretrained_checkpoint("rl_games", train_task_name) |
| 143 | if not resume_path: |
| 144 | print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.") |
| 145 | return |
| 146 | elif args_cli.checkpoint is None and not args_cli.use_last_checkpoint and DEFAULT_CHECKPOINT.exists(): |
| 147 | resume_path = str(DEFAULT_CHECKPOINT) |
| 148 | elif args_cli.checkpoint is None: |
| 149 | run_dir = agent_cfg["params"]["config"].get("full_experiment_name", ".*") |
| 150 | checkpoint_file = ".*" if args_cli.use_last_checkpoint else f"{agent_cfg['params']['config']['name']}.pth" |
| 151 | resume_path = get_checkpoint_path(log_root_path, run_dir, checkpoint_file, other_dirs=["nn"]) |
| 152 | else: |
| 153 | resume_path = retrieve_file_path(args_cli.checkpoint) |
| 154 | log_dir = os.path.dirname(os.path.dirname(resume_path)) |
| 155 | |
| 156 | rl_device = agent_cfg["params"]["config"]["device"] |
| 157 | clip_obs = agent_cfg["params"]["env"].get("clip_observations", math.inf) |
| 158 | clip_actions = agent_cfg["params"]["env"].get("clip_actions", math.inf) |
| 159 | obs_groups = agent_cfg["params"]["env"].get("obs_groups") |
| 160 | concate_obs_groups = agent_cfg["params"]["env"].get("concate_obs_groups", True) |
| 161 | |
| 162 | env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) |
| 163 | if isinstance(env.unwrapped, DirectMARLEnv): |
| 164 | env = multi_agent_to_single_agent(env) |
| 165 | if args_cli.video: |
| 166 | video_kwargs = { |
| 167 | "video_folder": os.path.join(log_root_path, log_dir, "videos", "play"), |
| 168 | "step_trigger": lambda step: step == 0, |
| 169 | "video_length": args_cli.video_length, |
| 170 | "disable_logger": True, |
| 171 | } |
| 172 | print("[INFO] Recording videos during play.") |
| 173 | print_dict(video_kwargs, nesting=4) |
| 174 | env = gym.wrappers.RecordVideo(env, **video_kwargs) |
| 175 | |
| 176 | base_env = env.unwrapped |
| 177 | handover_term = base_env.command_manager.get_term("handover") |
| 178 | current_command = _parse_handover_command(args_cli.initial_command) |
| 179 | handover_term.set_manual_command(current_command) |
| 180 | command_file_mtime = None |
| 181 | |
| 182 | env = RlGamesVecEnvWrapper(env, rl_device, clip_obs, clip_actions, obs_groups, concate_obs_groups) |
| 183 | vecenv.register("IsaacRlgWrapper", lambda config_name, num_actors, **kwargs: RlGamesGpuEnv(config_name, num_actors, **kwargs)) |
| 184 | env_configurations.register("rlgpu", {"vecenv_type": "IsaacRlgWrapper", "env_creator": lambda **kwargs: env}) |
no test coverage detected