Train with RSL-RL agent.
(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg)
| 114 | |
| 115 | @hydra_task_config(args_cli.task, args_cli.agent) |
| 116 | def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg): |
| 117 | """Train with RSL-RL agent.""" |
| 118 | # override configurations with non-hydra CLI arguments |
| 119 | agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli) |
| 120 | env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs |
| 121 | agent_cfg.max_iterations = ( |
| 122 | args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg.max_iterations |
| 123 | ) |
| 124 | |
| 125 | # set the environment seed |
| 126 | # note: certain randomizations occur in the environment initialization so we set the seed here |
| 127 | env_cfg.seed = agent_cfg.seed |
| 128 | env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device |
| 129 | # check for invalid combination of CPU device with distributed training |
| 130 | if args_cli.distributed and args_cli.device is not None and "cpu" in args_cli.device: |
| 131 | raise ValueError( |
| 132 | "Distributed training is not supported when using CPU device. " |
| 133 | "Please use GPU device (e.g., --device cuda) for distributed training." |
| 134 | ) |
| 135 | |
| 136 | # multi-gpu training configuration |
| 137 | if args_cli.distributed: |
| 138 | env_cfg.sim.device = f"cuda:{app_launcher.local_rank}" |
| 139 | agent_cfg.device = f"cuda:{app_launcher.local_rank}" |
| 140 | |
| 141 | # set seed to have diversity in different threads |
| 142 | seed = agent_cfg.seed + app_launcher.local_rank |
| 143 | env_cfg.seed = seed |
| 144 | agent_cfg.seed = seed |
| 145 | |
| 146 | # specify directory for logging experiments |
| 147 | log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name) |
| 148 | log_root_path = os.path.abspath(log_root_path) |
| 149 | print(f"[INFO] Logging experiment in directory: {log_root_path}") |
| 150 | # specify directory for logging runs: {time-stamp}_{run_name} |
| 151 | log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") |
| 152 | # The Ray Tune workflow extracts experiment name using the logging line below, hence, do not |
| 153 | # change it (see PR #2346, comment-2819298849) |
| 154 | print(f"Exact experiment name requested from command line: {log_dir}") |
| 155 | if agent_cfg.run_name: |
| 156 | log_dir += f"_{agent_cfg.run_name}" |
| 157 | log_dir = os.path.join(log_root_path, log_dir) |
| 158 | |
| 159 | # set the IO descriptors export flag if requested |
| 160 | if isinstance(env_cfg, ManagerBasedRLEnvCfg): |
| 161 | env_cfg.export_io_descriptors = args_cli.export_io_descriptors |
| 162 | else: |
| 163 | logger.warning( |
| 164 | "IO descriptors are only supported for manager based RL environments. No IO descriptors will be exported." |
| 165 | ) |
| 166 | |
| 167 | # set the log directory for the environment (works for all environment types) |
| 168 | env_cfg.log_dir = log_dir |
| 169 | |
| 170 | # create isaac environment |
| 171 | env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) |
| 172 | |
| 173 | # convert to single-agent instance if required by the RL algorithm |