()
| 239 | |
| 240 | |
| 241 | def main() -> None: |
| 242 | parser = argparse.ArgumentParser(description="Export Stage1/Stage2 checkpoint to ONNX.") |
| 243 | parser.add_argument("--checkpoint", type=str, required=True, help="Path to .pth (stage1) or .ckpt (stage2).") |
| 244 | parser.add_argument( |
| 245 | "--output", |
| 246 | type=str, |
| 247 | default="", |
| 248 | help="Output ONNX path. If empty, use policy_MMDDHH.onnx.", |
| 249 | ) |
| 250 | parser.add_argument("--stage", type=str, default="stage2", choices=["stage1", "stage2"]) |
| 251 | parser.add_argument("--config", type=str, default="", help="Optional config_*.yaml. Auto-resolved if empty.") |
| 252 | parser.add_argument("--obs_dim", type=int, default=DEFAULT_OBS_DIM) |
| 253 | parser.add_argument("--action_dim", type=int, default=DEFAULT_ACTION_DIM) |
| 254 | parser.add_argument("--prop_hist_len", type=int, default=DEFAULT_PROP_HIST_LEN) |
| 255 | parser.add_argument("--policy_rate", type=float, default=20.0, help="Policy control frequency in Hz.") |
| 256 | parser.add_argument("--chunk_size", type=int, default=1, help="Chunk size for action playback.") |
| 257 | parser.add_argument("--n_action_steps", type=int, default=1, help="Number of action steps per chunk.") |
| 258 | parser.add_argument("--opset", type=int, default=17) |
| 259 | parser.add_argument("--no_dynamic_batch", action="store_true", help="Disable dynamic batch axis.") |
| 260 | parser.add_argument( |
| 261 | "--meta_output", |
| 262 | type=str, |
| 263 | default="", |
| 264 | help="Optional deploy meta yaml path. Default: <output_without_ext>.deploy_meta.yaml", |
| 265 | ) |
| 266 | args = parser.parse_args() |
| 267 | |
| 268 | ckpt_path = Path(args.checkpoint).expanduser().resolve() |
| 269 | if not ckpt_path.exists(): |
| 270 | raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}") |
| 271 | |
| 272 | run_dir = _find_run_dir_for_checkpoint(ckpt_path) |
| 273 | onnx_dir = Path("outputs") / "hora" / "revo3_right" / "onnx" |
| 274 | output_name = args.output.strip() if isinstance(args.output, str) else "" |
| 275 | if not output_name: |
| 276 | # Default: outputs/hora/revo3_right/onnx/policy_MMDDHH.onnx |
| 277 | out_path = (onnx_dir / f"policy_{datetime.datetime.now().strftime('%m%d%H')}.onnx").resolve() |
| 278 | else: |
| 279 | out_arg_path = Path(output_name).expanduser() |
| 280 | if out_arg_path.is_absolute() or out_arg_path.parent != Path("."): |
| 281 | out_path = out_arg_path.resolve() |
| 282 | else: |
| 283 | # If only a filename is given, still place it under outputs/hora/revo3_right/onnx/ |
| 284 | out_path = (onnx_dir / out_arg_path.name).resolve() |
| 285 | out_path.parent.mkdir(parents=True, exist_ok=True) |
| 286 | |
| 287 | cfg_path = Path(args.config).expanduser().resolve() if args.config else _find_config_for_checkpoint(ckpt_path) |
| 288 | cfg = _load_config(cfg_path) |
| 289 | if cfg_path is None: |
| 290 | print("[WARN] config_*.yaml not found near checkpoint; using script defaults.") |
| 291 | else: |
| 292 | print(f"[INFO] Using config: {cfg_path}") |
| 293 | |
| 294 | obs_dim = int(args.obs_dim) |
| 295 | actions_num = int(args.action_dim) |
| 296 | prop_hist_len = int(args.prop_hist_len) |
| 297 | scale_keys = _resolve_scale_keys_from_config(cfg) |
| 298 | obs_per_step = obs_dim // 3 |
no test coverage detected