(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: dict)
| 173 | |
| 174 | @hydra_task_config(args_cli.task, args_cli.agent) |
| 175 | def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: dict): |
| 176 | env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs |
| 177 | env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device |
| 178 | |
| 179 | if args_cli.seed == -1: |
| 180 | args_cli.seed = random.randint(0, 10000) |
| 181 | |
| 182 | agent_cfg["params"]["seed"] = args_cli.seed if args_cli.seed is not None else agent_cfg["params"]["seed"] |
| 183 | agent_cfg["params"]["config"]["max_epochs"] = ( |
| 184 | args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg["params"]["config"]["max_epochs"] |
| 185 | ) |
| 186 | if args_cli.checkpoint is not None: |
| 187 | resume_path = retrieve_file_path(args_cli.checkpoint) |
| 188 | agent_cfg["params"]["load_checkpoint"] = True |
| 189 | agent_cfg["params"]["load_path"] = resume_path |
| 190 | print(f"[INFO] Loading model checkpoint from: {agent_cfg['params']['load_path']}") |
| 191 | train_sigma = float(args_cli.sigma) if args_cli.sigma is not None else None |
| 192 | |
| 193 | if args_cli.distributed: |
| 194 | agent_cfg["params"]["seed"] += app_launcher.global_rank |
| 195 | agent_cfg["params"]["config"]["device"] = f"cuda:{app_launcher.local_rank}" |
| 196 | agent_cfg["params"]["config"]["device_name"] = f"cuda:{app_launcher.local_rank}" |
| 197 | agent_cfg["params"]["config"]["multi_gpu"] = True |
| 198 | env_cfg.sim.device = f"cuda:{app_launcher.local_rank}" |
| 199 | |
| 200 | env_cfg.seed = agent_cfg["params"]["seed"] |
| 201 | |
| 202 | config_name = agent_cfg["params"]["config"]["name"] |
| 203 | log_root_path = os.path.abspath(os.path.join("logs", "rl_games", config_name)) |
| 204 | print(f"[INFO] Logging experiment in directory: {log_root_path}") |
| 205 | log_dir = agent_cfg["params"]["config"].get("full_experiment_name", datetime.now().strftime("%Y-%m-%d_%H-%M-%S")) |
| 206 | agent_cfg["params"]["config"]["train_dir"] = log_root_path |
| 207 | agent_cfg["params"]["config"]["full_experiment_name"] = log_dir |
| 208 | |
| 209 | dump_yaml(os.path.join(log_root_path, log_dir, "params", "env.yaml"), env_cfg) |
| 210 | dump_yaml(os.path.join(log_root_path, log_dir, "params", "agent.yaml"), agent_cfg) |
| 211 | dump_pickle(os.path.join(log_root_path, log_dir, "params", "env.pkl"), env_cfg) |
| 212 | dump_pickle(os.path.join(log_root_path, log_dir, "params", "agent.pkl"), agent_cfg) |
| 213 | |
| 214 | rl_device = agent_cfg["params"]["config"]["device"] |
| 215 | clip_obs = agent_cfg["params"]["env"].get("clip_observations", math.inf) |
| 216 | clip_actions = agent_cfg["params"]["env"].get("clip_actions", math.inf) |
| 217 | obs_groups = agent_cfg["params"]["env"].get("obs_groups") |
| 218 | concate_obs_groups = agent_cfg["params"]["env"].get("concate_obs_groups", True) |
| 219 | |
| 220 | if isinstance(env_cfg, ManagerBasedRLEnvCfg): |
| 221 | env_cfg.export_io_descriptors = args_cli.export_io_descriptors |
| 222 | env_cfg.io_descriptors_output_dir = os.path.join(log_root_path, log_dir) |
| 223 | else: |
| 224 | omni.log.warn("IO descriptors are only supported for manager based RL environments.") |
| 225 | |
| 226 | env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) |
| 227 | if isinstance(env.unwrapped, DirectMARLEnv): |
| 228 | env = multi_agent_to_single_agent(env) |
| 229 | |
| 230 | if args_cli.video: |
| 231 | video_kwargs = { |
| 232 | "video_folder": os.path.join(log_root_path, log_dir, "videos", "train"), |
no test coverage detected