Main function.
()
| 317 | |
| 318 | |
| 319 | def main(): |
| 320 | """Main function.""" |
| 321 | # Create output directory |
| 322 | os.makedirs(args_cli.output_dir, exist_ok=True) |
| 323 | |
| 324 | # Load the pkl file |
| 325 | print(f"Loading pkl file: {args_cli.input_file}") |
| 326 | motions_data = joblib.load(args_cli.input_file) |
| 327 | print(f"Loaded {len(motions_data)} motions") |
| 328 | |
| 329 | # Limit number of motions if specified |
| 330 | if args_cli.max_motions: |
| 331 | motions_data = dict(list(motions_data.items())[: args_cli.max_motions]) |
| 332 | print(f"Limited to {len(motions_data)} motions for processing") |
| 333 | |
| 334 | # Load kit helper |
| 335 | sim_cfg = sim_utils.SimulationCfg(device=args_cli.device) |
| 336 | sim_cfg.dt = 1.0 / args_cli.output_fps |
| 337 | sim = SimulationContext(sim_cfg) |
| 338 | # Design scene |
| 339 | scene_cfg = ReplayMotionsSceneCfg(num_envs=1, env_spacing=2.0) |
| 340 | scene = InteractiveScene(scene_cfg) |
| 341 | # Play the simulator |
| 342 | sim.reset() |
| 343 | # Now we are ready! |
| 344 | print("[INFO]: Setup complete...") |
| 345 | |
| 346 | # Process each motion |
| 347 | joint_names = [ |
| 348 | "left_hip_pitch_joint", |
| 349 | "left_hip_roll_joint", |
| 350 | "left_hip_yaw_joint", |
| 351 | "left_knee_joint", |
| 352 | "left_ankle_pitch_joint", |
| 353 | "left_ankle_roll_joint", |
| 354 | "right_hip_pitch_joint", |
| 355 | "right_hip_roll_joint", |
| 356 | "right_hip_yaw_joint", |
| 357 | "right_knee_joint", |
| 358 | "right_ankle_pitch_joint", |
| 359 | "right_ankle_roll_joint", |
| 360 | "waist_yaw_joint", |
| 361 | "waist_roll_joint", |
| 362 | "waist_pitch_joint", |
| 363 | "left_shoulder_pitch_joint", |
| 364 | "left_shoulder_roll_joint", |
| 365 | "left_shoulder_yaw_joint", |
| 366 | "left_elbow_joint", |
| 367 | "left_wrist_roll_joint", |
| 368 | "left_wrist_pitch_joint", |
| 369 | "left_wrist_yaw_joint", |
| 370 | "right_shoulder_pitch_joint", |
| 371 | "right_shoulder_roll_joint", |
| 372 | "right_shoulder_yaw_joint", |
| 373 | "right_elbow_joint", |
| 374 | "right_wrist_roll_joint", |
| 375 | "right_wrist_pitch_joint", |
| 376 | "right_wrist_yaw_joint", |
no test coverage detected