()
| 178 | |
| 179 | |
| 180 | def main(): |
| 181 | if args.test and not args.checkpoint: |
| 182 | raise ValueError('--test requires --checkpoint') |
| 183 | if args.algo == 'ProprioAdapt' and not args.checkpoint: |
| 184 | raise ValueError('ProprioAdapt training requires --checkpoint') |
| 185 | |
| 186 | set_np_formatting() |
| 187 | seed = set_seed(args.seed) |
| 188 | full_config = _build_full_config(seed) |
| 189 | |
| 190 | cprint('Start Building the Environment', 'green', attrs=['bold']) |
| 191 | env_cfg = _build_env_cfg(seed) |
| 192 | if args.algo == 'ProprioAdapt': |
| 193 | env_cfg.enable_contact_in_obs = False # Stage2: actor sees zero contact, adapt_tconv still sees contact history |
| 194 | if args.test: |
| 195 | env_cfg.gravity_curriculum = False |
| 196 | env_cfg.sim.gravity = (0.0, 0.0, -9.81) # full gravity for test/play |
| 197 | env = Revo3HandHoraEnv( |
| 198 | cfg=env_cfg, |
| 199 | render_mode=None if getattr(args, 'headless', False) else 'human', |
| 200 | ) |
| 201 | env = HoraCompatWrapper(env) |
| 202 | |
| 203 | # Output to Stage1's run directory |
| 204 | if args.algo == 'ProprioAdapt' and not _is_stage2_checkpoint(args.checkpoint): |
| 205 | output_dif = os.path.dirname(os.path.dirname(args.checkpoint)) |
| 206 | else: |
| 207 | output_dif = os.path.join('outputs', 'hora', 'revo3_right', args.output_name) |
| 208 | os.makedirs(output_dif, exist_ok=True) |
| 209 | algo_name = str(full_config.train.algo) |
| 210 | if algo_name not in _ALGO_MAP: |
| 211 | raise ValueError(f"Unsupported algo: {algo_name}. Available: {list(_ALGO_MAP.keys())}") |
| 212 | agent = _ALGO_MAP[algo_name](env, output_dif, full_config=full_config) |
| 213 | |
| 214 | if args.test: |
| 215 | agent.restore_test(full_config.train.load_path) |
| 216 | agent.test() |
| 217 | else: |
| 218 | best_ckpt_path = os.path.join( |
| 219 | output_dif, |
| 220 | 'stage1_nn' if full_config.train.algo == 'PPO' else 'stage2_nn', |
| 221 | 'best.pth' if full_config.train.algo == 'PPO' else 'model_best.ckpt', |
| 222 | ) |
| 223 | if os.path.exists(best_ckpt_path): |
| 224 | if args.force_overwrite: |
| 225 | print(f"[INFO] --force_overwrite enabled, continue and overwrite in {output_dif}", flush=True) |
| 226 | else: |
| 227 | user_input = input( |
| 228 | f'are you intentionally going to overwrite files in {output_dif}, type yes to continue \n' |
| 229 | ) |
| 230 | if user_input != 'yes': |
| 231 | return |
| 232 | |
| 233 | _attach_env_runtime_to_config(full_config, env_cfg) |
| 234 | _save_run_metadata(output_dif, full_config) |
| 235 | agent.restore_train(full_config.train.load_path) |
| 236 | agent.train() |
| 237 |
no test coverage detected