From 501455d3916661dae3391ee66cbae8995a3cfd57 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 5 Sep 2026 17:50:04 +0000 Subject: [PATCH 1/2] feat(motion): add fixed-scene expert trajectory generation Group solvers, planners, workspace and augmentation under sim.motion. Add fixed-scene preparation, bounded generation and confirmed LeRobot persistence, plus physically validated parallel cube PickUp collection and video examples. BREAKING CHANGE: import robot motion APIs from embodichain.lab.sim.motion; previous sim.solvers, sim.planners and sim.workspace paths are removed. --- .agents/skills/add-solver/SKILL.md | 48 +- .agents/skills/add-test/SKILL.md | 2 +- .../review-pr/references/review-matrix.md | 7 +- AGENTS.md | 18 +- agent_context/MAP.yaml | 93 +- .../topics/atomic-actions/atomic-actions.md | 30 + .../topics/env-framework/env-framework.md | 62 ++ agent_context/topics/ik-solvers/ik-solvers.md | 15 +- .../topics/motion-planning/motion-planning.md | 136 ++- .../topics/robot-system/robot-system.md | 20 + .../simulation-system/simulation-system.md | 40 +- ...e_expert_trajectory_augmentation_design.md | 578 ++++++++++++ ...ectory_augmentation_implementation_plan.md | 351 +++++++ .../embodichain/embodichain.lab.gym.envs.rst | 37 + .../embodichain/embodichain.lab.gym.rst | 7 +- .../embodichain/embodichain.lab.rst | 12 +- ...> embodichain.lab.sim.motion.planners.rst} | 8 +- .../embodichain.lab.sim.motion.rst | 45 + ...=> embodichain.lab.sim.motion.solvers.rst} | 8 +- ...lab.sim.motion.trajectory_augmentation.rst | 246 +++++ ...ichain.lab.sim.motion.workspace.caches.rst | 7 + ...chain.lab.sim.motion.workspace.configs.rst | 7 + ...n.lab.sim.motion.workspace.constraints.rst | 7 + ...chain.lab.sim.motion.workspace.metrics.rst | 7 + ... embodichain.lab.sim.motion.workspace.rst} | 14 +- ...hain.lab.sim.motion.workspace.samplers.rst | 7 + ...n.lab.sim.motion.workspace.visualizers.rst | 7 + .../embodichain/embodichain.lab.sim.rst | 34 +- .../embodichain.lab.sim.sim_manager.rst | 5 + .../embodichain.lab.sim.workspace.caches.rst | 7 - .../embodichain.lab.sim.workspace.configs.rst | 7 - ...odichain.lab.sim.workspace.constraints.rst | 7 - .../embodichain.lab.sim.workspace.metrics.rst | 7 - ...embodichain.lab.sim.workspace.samplers.rst | 7 - ...odichain.lab.sim.workspace.visualizers.rst | 7 - .../embodichain.lab.trajectory_generation.rst | 390 ++++++++ .../embodichain/embodichain.toolkits.rst | 5 +- docs/source/api_reference/public_api.rst | 776 +++++++-------- .../features/workspace_analyzer/caches.md | 10 +- .../source/features/workspace_analyzer/cli.md | 4 +- .../features/workspace_analyzer/configs.md | 14 +- .../workspace_analyzer/constraints.md | 34 +- .../features/workspace_analyzer/index.rst | 2 +- .../features/workspace_analyzer/metrics.md | 14 +- .../features/workspace_analyzer/runtime.md | 6 +- .../features/workspace_analyzer/samplers.md | 8 +- .../workspace_analyzer/visualizers.md | 14 +- .../workspace_analyzer/workspace_analyzer.md | 10 +- docs/source/guides/add_robot.rst | 2 +- docs/source/index.rst | 1 + docs/source/overview/gym/index.rst | 2 + .../overview/sim/atomic_actions/index.md | 2 +- docs/source/overview/sim/index.rst | 31 +- docs/source/overview/sim/motion/index.rst | 47 + .../{ => motion}/planners/curobo_planner.md | 23 +- .../sim/{ => motion}/planners/index.rst | 0 .../{ => motion}/planners/motion_generator.md | 12 +- .../{ => motion}/planners/neural_planner.md | 6 +- .../{ => motion}/planners/toppra_planner.md | 6 +- .../planners/trajectory_sample_method.md | 0 .../solvers/differential_solver.md | 4 +- .../sim/{ => motion}/solvers/index.rst | 0 .../{ => motion}/solvers/neural_ik_solver.md | 2 +- .../sim/{ => motion}/solvers/opw_solver.md | 2 +- .../sim/{ => motion}/solvers/pink_solver.md | 4 +- .../{ => motion}/solvers/pinocchio_solver.md | 4 +- .../{ => motion}/solvers/pytorch_solver.md | 4 +- .../sim/{ => motion}/solvers/srs_solver.md | 2 +- .../sim/{ => motion}/solvers/ur_solver.md | 2 +- docs/source/overview/sim/sim_robot.md | 2 +- .../overview/task_program/scene_registry.md | 4 +- .../source/overview/trajectory_generation.rst | 566 +++++++++++ docs/source/quick_start/install.md | 2 +- docs/source/resources/robot/franka_panda.md | 2 +- docs/source/resources/robot/ur_robot.md | 2 +- docs/source/tutorial/atomic_actions.rst | 2 +- docs/source/tutorial/motion_gen.rst | 10 +- docs/source/tutorial/solver.rst | 2 +- embodichain/cli/workspace_cache.py | 2 +- .../envs/action_bank/configurable_action.py | 2 +- embodichain/lab/gym/envs/base_env.py | 133 ++- embodichain/lab/gym/envs/demo.py | 108 ++- embodichain/lab/gym/envs/embodied_env.py | 151 ++- embodichain/lab/scripts/analyze_workspace.py | 14 +- embodichain/lab/sim/atomic_actions/core.py | 19 +- embodichain/lab/sim/atomic_actions/engine.py | 27 +- .../lab/sim/atomic_actions/execution.py | 19 +- embodichain/lab/sim/atomic_actions/plans.py | 2 +- .../lab/sim/atomic_actions/policies.py | 6 +- .../sim/atomic_actions/primitives/_helpers.py | 2 +- .../sim/atomic_actions/primitives/pick_up.py | 74 +- embodichain/lab/sim/atomic_actions/runtime.py | 2 +- .../lab/sim/atomic_actions/sim_adapter.py | 2 +- .../lab/sim/atomic_actions/trajectory_ops.py | 2 +- embodichain/lab/sim/cfg.py | 6 +- embodichain/lab/sim/motion/__init__.py | 43 + .../lab/sim/{ => motion}/planners/__init__.py | 2 + .../sim/{ => motion}/planners/base_planner.py | 0 .../planners/curobo/curobo_planner.py | 97 +- .../planners/curobo/curobo_yaml.py | 4 +- .../{ => motion}/planners/motion_generator.py | 2 +- .../{ => motion}/planners/neural_planner.py | 4 +- .../{ => motion}/planners/toppra_planner.py | 4 +- .../lab/sim/{ => motion}/planners/utils.py | 0 .../lab/sim/{ => motion}/solvers/__init__.py | 23 + .../sim/{ => motion}/solvers/base_solver.py | 2 +- .../solvers/differential_solver.py | 2 +- .../{ => motion}/solvers/neural_ik_solver.py | 4 +- .../solvers/null_space_posture_task.py | 0 .../sim/{ => motion}/solvers/opw_solver.py | 2 +- .../sim/{ => motion}/solvers/pink_solver.py | 8 +- .../{ => motion}/solvers/pinocchio_solver.py | 2 +- .../{ => motion}/solvers/pytorch_solver.py | 4 +- .../{ => motion}/solvers/qpos_seed_sampler.py | 0 .../sim/{ => motion}/solvers/srs_solver.py | 2 +- .../lab/sim/{ => motion}/solvers/ur_solver.py | 2 +- .../trajectory_augmentation/__init__.py | 70 ++ .../sim/motion/trajectory_augmentation/cfg.py | 436 +++++++++ .../trajectory_augmentation/contracts.py | 582 ++++++++++++ .../trajectory_augmentation/coverage.py | 286 ++++++ .../trajectory_augmentation/operators.py | 319 +++++++ .../motion/trajectory_augmentation/session.py | 637 +++++++++++++ .../sim/{ => motion}/workspace/__init__.py | 0 .../sim/{ => motion}/workspace/analyzer.py | 16 +- .../{ => motion}/workspace/caches/__init__.py | 10 +- .../workspace/caches/base_cache.py | 0 .../workspace/caches/cache_manager.py | 8 +- .../workspace/caches/cache_utils.py | 0 .../workspace/caches/disk_cache.py | 2 +- .../workspace/caches/memory_cache.py | 2 +- .../workspace/caches/results_cache.py | 0 .../lab/sim/{ => motion}/workspace/cfg.py | 0 .../workspace/configs/__init__.py | 10 +- .../workspace/configs/cache_config.py | 0 .../workspace/configs/dimension_constraint.py | 0 .../workspace/configs/metric_config.py | 0 .../workspace/configs/sampling_config.py | 0 .../workspace/configs/visualization_config.py | 0 .../workspace/constraints/__init__.py | 0 .../workspace/constraints/base_constraint.py | 0 .../constraints/workspace_constraint.py | 12 +- .../workspace/metrics/__init__.py | 8 +- .../workspace/metrics/base_metric.py | 0 .../workspace/metrics/density_metric.py | 4 +- .../metrics/manipulability_metric.py | 4 +- .../workspace/metrics/reachability_metric.py | 4 +- .../lab/sim/{ => motion}/workspace/runtime.py | 0 .../workspace/samplers/__init__.py | 16 +- .../workspace/samplers/base_sampler.py | 0 .../workspace/samplers/gaussian_sampler.py | 4 +- .../workspace/samplers/halton_sampler.py | 4 +- .../workspace/samplers/importance_sampler.py | 4 +- .../workspace/samplers/iniform_sampler.py | 4 +- .../workspace/samplers/lhs_sampler.py | 4 +- .../workspace/samplers/random_sampler.py | 4 +- .../workspace/samplers/sampler_factory.py | 18 +- .../workspace/samplers/sobol_sampler.py | 4 +- .../workspace/visualizers/__init__.py | 14 +- .../workspace/visualizers/axis_visualizer.py | 4 +- .../workspace/visualizers/base_visualizer.py | 2 +- .../visualizers/point_cloud_visualizer.py | 4 +- .../visualizers/sphere_visualizer.py | 4 +- .../visualizers/visualizer_factory.py | 12 +- .../workspace/visualizers/voxel_visualizer.py | 4 +- embodichain/lab/sim/objects/articulation.py | 2 +- embodichain/lab/sim/objects/gizmo.py | 2 +- embodichain/lab/sim/objects/robot.py | 4 +- embodichain/lab/sim/robots/cobotmagic.py | 2 +- embodichain/lab/sim/robots/dexforce_w1/cfg.py | 2 +- embodichain/lab/sim/robots/dual_arm.py | 2 +- embodichain/lab/sim/robots/franka_panda.py | 4 +- embodichain/lab/sim/robots/ur_robot.py | 4 +- embodichain/lab/sim/sim_manager.py | 17 + embodichain/lab/sim/utility/action_utils.py | 2 +- .../lab/sim/utility/atom_action_utils.py | 2 +- embodichain/lab/sim/utility/cfg_utils.py | 2 +- .../integrations/simulation/environment.py | 2 +- .../simulation/parallel_safety.py | 2 +- .../lab/task_program/semantics/scene.py | 2 +- .../lab/trajectory_generation/__init__.py | 19 + .../lab/trajectory_generation/execution.py | 891 ++++++++++++++++++ .../trajectory_generation/initial_state.py | 398 ++++++++ .../integrations/__init__.py | 19 + .../integrations/_collision.py | 189 ++++ .../integrations/atomic.py | 163 ++++ .../integrations/contact.py | 661 +++++++++++++ .../integrations/planning.py | 614 ++++++++++++ .../trajectory_generation/integrations/sim.py | 433 +++++++++ .../lab/trajectory_generation/runner.py | 743 +++++++++++++++ .../lab/trajectory_generation/sinks.py | 553 +++++++++++ .../tableware/blocks_ranking_rgb.py | 2 +- .../tableware/stack_blocks_two.py | 2 +- examples/sim/gizmo/gizmo_robot.py | 2 +- .../{ => motion}/planners/curobo_planner.py | 12 +- .../{ => motion}/planners/neural_planner.py | 8 +- .../solvers/differential_solver.py | 0 .../{ => motion}/solvers/neural_ik_solver.py | 2 +- .../sim/{ => motion}/solvers/opw_solver.py | 0 .../sim/{ => motion}/solvers/pink_solver.py | 0 .../{ => motion}/solvers/pinocchio_solver.py | 0 .../{ => motion}/solvers/pytorch_solver.py | 0 .../sim/{ => motion}/solvers/srs_solver.py | 0 .../motion/trajectory_generation/README.md | 97 ++ .../cube_grasp_parallel.py | 544 +++++++++++ .../cube_pickup_collection.py | 416 ++++++++ .../trajectory_generation/free_motion.py | 445 +++++++++ .../workspace/analyze_cartesian_workspace.py | 4 +- .../workspace/analyze_joint_workspace.py | 4 +- .../workspace/analyze_plane_workspace.py | 4 +- pyproject.toml | 5 + .../move_end_effector_benchmark.py | 4 +- .../move_held_object_benchmark.py | 4 +- .../atomic_action/move_joints_benchmark.py | 4 +- .../atomic_action/pickup_benchmark.py | 4 +- .../atomic_action/place_benchmark.py | 4 +- .../curobo_extraction/run_benchmark.py | 4 +- .../motion_generation/BENCHMARK_DESIGN.md | 2 +- .../motion_generation/metrics/trajectory.py | 2 +- .../motion_generation/planners/base.py | 2 +- .../motion_generation/planners/curobo.py | 2 +- .../planners/ik_interpolate.py | 2 +- .../motion_generation/planners/neural.py | 2 +- .../motion_generation/planners/toppra.py | 4 +- scripts/benchmark/motion_generation/runner.py | 2 +- .../kinematic_solver/run_benchmark.py | 9 +- .../robotics/kinematic_solver/srs_solver.py | 2 +- .../benchmark_workspace_analyzer.py | 14 +- .../dynamic_obstacle_recovery.py | 4 +- scripts/tutorials/atomic_action/pour.py | 5 +- .../tutorials/atomic_action/scenario_utils.py | 2 +- .../tutorials/atomic_action/tutorial_utils.py | 2 +- scripts/tutorials/grasp/grasp_generator.py | 2 +- scripts/tutorials/sim/gizmo_robot.py | 2 +- scripts/tutorials/sim/motion_generator.py | 4 +- scripts/tutorials/sim/open_drawer.py | 2 +- scripts/tutorials/visualization/README.md | 4 +- .../test_motion_generation_benchmark.py | 4 +- .../managers/test_workspace_randomization.py | 2 +- .../test_simulation_environment.py | 2 +- .../test_simulation_parallel_safety.py | 2 +- tests/gym/envs/test_demo.py | 145 ++- .../gym/envs/test_fixed_scene_preparation.py | 523 ++++++++++ .../lab/task_program/semantics/test_scene.py | 2 +- .../test_scene_curobo_integration.py | 2 +- .../test_semantic_executor_curobo_gpu.py | 7 +- tests/lab/trajectory_generation/__init__.py | 19 + .../test_atomic_source.py | 103 ++ .../lab/trajectory_generation/test_contact.py | 268 ++++++ .../test_cube_grasp_parallel.py | 132 +++ .../test_episode_sinks.py | 331 +++++++ .../trajectory_generation/test_execution.py | 684 ++++++++++++++ .../test_initial_state_host.py | 266 ++++++ .../test_pickup_collection.py | 176 ++++ .../trajectory_generation/test_planning.py | 570 +++++++++++ .../lab/trajectory_generation/test_runner.py | 506 ++++++++++ .../trajectory_generation/test_runner_real.py | 101 ++ .../test_session_sink_integration.py | 111 +++ .../test_sim_initial_state.py | 463 +++++++++ tests/sim/atomic_actions/test_actions.py | 250 ++++- tests/sim/atomic_actions/test_core.py | 2 +- .../test_curobo_motion_strategy_e2e.py | 4 +- .../test_endpoint_runtime_e2e.py | 2 +- .../sim/atomic_actions/test_engine_per_env.py | 233 ++++- .../test_motion_strategy_e2e.py | 6 +- .../sim/atomic_actions/test_trajectory_ops.py | 2 +- .../sim/atomic_actions/test_tutorial_utils.py | 6 +- .../planners/test_base_planner.py | 2 +- .../planners/test_curobo_integration.py | 4 +- .../planners/test_curobo_planner.py | 145 ++- .../planners/test_motion_generator.py | 4 +- .../planners/test_motion_generator_batched.py | 11 +- .../planners/test_neural_batched.py | 12 +- .../planners/test_neural_planner.py | 6 +- .../planners/test_plan_state_batched.py | 11 +- .../planners/test_toppra_batched.py | 45 +- .../planners/test_toppra_planner.py | 4 +- .../{ => motion}/solvers/test_base_solver.py | 8 +- .../solvers/test_differential_solver.py | 0 .../solvers/test_neural_ik_solver.py | 4 +- .../{ => motion}/solvers/test_opw_solver.py | 0 .../{ => motion}/solvers/test_pink_solver.py | 6 +- .../solvers/test_pinocchio_solver.py | 0 .../solvers/test_pytorch_solver.py | 0 .../{ => motion}/solvers/test_srs_solver.py | 2 +- .../{ => motion}/solvers/test_ur_solver.py | 2 +- tests/sim/motion/test_motion_imports.py | 67 ++ .../trajectory_augmentation/__init__.py | 17 + .../trajectory_augmentation/test_cfg.py | 175 ++++ .../trajectory_augmentation/test_contracts.py | 301 ++++++ .../trajectory_augmentation/test_coverage.py | 144 +++ .../trajectory_augmentation/test_imports.py | 33 + .../trajectory_augmentation/test_operators.py | 233 +++++ .../trajectory_augmentation/test_session.py | 572 +++++++++++ tests/sim/{ => motion}/workspace/__init__.py | 0 .../{ => motion}/workspace/test_analyzer.py | 10 +- .../sim/{ => motion}/workspace/test_cache.py | 28 +- .../{ => motion}/workspace/test_runtime.py | 2 +- tests/sim/objects/test_articulation.py | 11 + tests/sim/objects/test_dual_arm.py | 2 +- tests/sim/objects/test_gizmo.py | 2 +- tests/sim/objects/test_robot_cfg.py | 8 +- tests/toolkits/test_grasp_pose_generator.py | 2 +- 302 files changed, 19339 insertions(+), 1062 deletions(-) create mode 100644 docs/design/fixed_scene_expert_trajectory_augmentation_design.md create mode 100644 docs/design/fixed_scene_expert_trajectory_augmentation_implementation_plan.md rename docs/source/api_reference/embodichain/{embodichain.lab.sim.planners.rst => embodichain.lab.sim.motion.planners.rst} (88%) create mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.motion.rst rename docs/source/api_reference/embodichain/{embodichain.lab.sim.solvers.rst => embodichain.lab.sim.motion.solvers.rst} (94%) create mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.motion.trajectory_augmentation.rst create mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.caches.rst create mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.configs.rst create mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.constraints.rst create mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.metrics.rst rename docs/source/api_reference/embodichain/{embodichain.lab.sim.workspace.rst => embodichain.lab.sim.motion.workspace.rst} (60%) create mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.samplers.rst create mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.visualizers.rst delete mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.caches.rst delete mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.configs.rst delete mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.constraints.rst delete mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.metrics.rst delete mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.samplers.rst delete mode 100644 docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.visualizers.rst create mode 100644 docs/source/api_reference/embodichain/embodichain.lab.trajectory_generation.rst create mode 100644 docs/source/overview/sim/motion/index.rst rename docs/source/overview/sim/{ => motion}/planners/curobo_planner.md (94%) rename docs/source/overview/sim/{ => motion}/planners/index.rst (100%) rename docs/source/overview/sim/{ => motion}/planners/motion_generator.md (93%) rename docs/source/overview/sim/{ => motion}/planners/neural_planner.md (87%) rename docs/source/overview/sim/{ => motion}/planners/toppra_planner.md (90%) rename docs/source/overview/sim/{ => motion}/planners/trajectory_sample_method.md (100%) rename docs/source/overview/sim/{ => motion}/solvers/differential_solver.md (95%) rename docs/source/overview/sim/{ => motion}/solvers/index.rst (100%) rename docs/source/overview/sim/{ => motion}/solvers/neural_ik_solver.md (96%) rename docs/source/overview/sim/{ => motion}/solvers/opw_solver.md (98%) rename docs/source/overview/sim/{ => motion}/solvers/pink_solver.md (96%) rename docs/source/overview/sim/{ => motion}/solvers/pinocchio_solver.md (95%) rename docs/source/overview/sim/{ => motion}/solvers/pytorch_solver.md (96%) rename docs/source/overview/sim/{ => motion}/solvers/srs_solver.md (98%) rename docs/source/overview/sim/{ => motion}/solvers/ur_solver.md (99%) create mode 100644 docs/source/overview/trajectory_generation.rst create mode 100644 embodichain/lab/sim/motion/__init__.py rename embodichain/lab/sim/{ => motion}/planners/__init__.py (97%) rename embodichain/lab/sim/{ => motion}/planners/base_planner.py (100%) rename embodichain/lab/sim/{ => motion}/planners/curobo/curobo_planner.py (96%) rename embodichain/lab/sim/{ => motion}/planners/curobo/curobo_yaml.py (99%) rename embodichain/lab/sim/{ => motion}/planners/motion_generator.py (99%) rename embodichain/lab/sim/{ => motion}/planners/neural_planner.py (99%) rename embodichain/lab/sim/{ => motion}/planners/toppra_planner.py (99%) rename embodichain/lab/sim/{ => motion}/planners/utils.py (100%) rename embodichain/lab/sim/{ => motion}/solvers/__init__.py (77%) rename embodichain/lab/sim/{ => motion}/solvers/base_solver.py (99%) rename embodichain/lab/sim/{ => motion}/solvers/differential_solver.py (99%) rename embodichain/lab/sim/{ => motion}/solvers/neural_ik_solver.py (98%) rename embodichain/lab/sim/{ => motion}/solvers/null_space_posture_task.py (100%) rename embodichain/lab/sim/{ => motion}/solvers/opw_solver.py (99%) rename embodichain/lab/sim/{ => motion}/solvers/pink_solver.py (98%) rename embodichain/lab/sim/{ => motion}/solvers/pinocchio_solver.py (99%) rename embodichain/lab/sim/{ => motion}/solvers/pytorch_solver.py (99%) rename embodichain/lab/sim/{ => motion}/solvers/qpos_seed_sampler.py (100%) rename embodichain/lab/sim/{ => motion}/solvers/srs_solver.py (99%) rename embodichain/lab/sim/{ => motion}/solvers/ur_solver.py (99%) create mode 100644 embodichain/lab/sim/motion/trajectory_augmentation/__init__.py create mode 100644 embodichain/lab/sim/motion/trajectory_augmentation/cfg.py create mode 100644 embodichain/lab/sim/motion/trajectory_augmentation/contracts.py create mode 100644 embodichain/lab/sim/motion/trajectory_augmentation/coverage.py create mode 100644 embodichain/lab/sim/motion/trajectory_augmentation/operators.py create mode 100644 embodichain/lab/sim/motion/trajectory_augmentation/session.py rename embodichain/lab/sim/{ => motion}/workspace/__init__.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/analyzer.py (99%) rename embodichain/lab/sim/{ => motion}/workspace/caches/__init__.py (77%) rename embodichain/lab/sim/{ => motion}/workspace/caches/base_cache.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/caches/cache_manager.py (91%) rename embodichain/lab/sim/{ => motion}/workspace/caches/cache_utils.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/caches/disk_cache.py (98%) rename embodichain/lab/sim/{ => motion}/workspace/caches/memory_cache.py (96%) rename embodichain/lab/sim/{ => motion}/workspace/caches/results_cache.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/cfg.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/configs/__init__.py (81%) rename embodichain/lab/sim/{ => motion}/workspace/configs/cache_config.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/configs/dimension_constraint.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/configs/metric_config.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/configs/sampling_config.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/configs/visualization_config.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/constraints/__init__.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/constraints/base_constraint.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/constraints/workspace_constraint.py (95%) rename embodichain/lab/sim/{ => motion}/workspace/metrics/__init__.py (78%) rename embodichain/lab/sim/{ => motion}/workspace/metrics/base_metric.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/metrics/density_metric.py (97%) rename embodichain/lab/sim/{ => motion}/workspace/metrics/manipulability_metric.py (97%) rename embodichain/lab/sim/{ => motion}/workspace/metrics/reachability_metric.py (96%) rename embodichain/lab/sim/{ => motion}/workspace/runtime.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/samplers/__init__.py (75%) rename embodichain/lab/sim/{ => motion}/workspace/samplers/base_sampler.py (100%) rename embodichain/lab/sim/{ => motion}/workspace/samplers/gaussian_sampler.py (98%) rename embodichain/lab/sim/{ => motion}/workspace/samplers/halton_sampler.py (98%) rename embodichain/lab/sim/{ => motion}/workspace/samplers/importance_sampler.py (98%) rename embodichain/lab/sim/{ => motion}/workspace/samplers/iniform_sampler.py (97%) rename embodichain/lab/sim/{ => motion}/workspace/samplers/lhs_sampler.py (98%) rename embodichain/lab/sim/{ => motion}/workspace/samplers/random_sampler.py (96%) rename embodichain/lab/sim/{ => motion}/workspace/samplers/sampler_factory.py (92%) rename embodichain/lab/sim/{ => motion}/workspace/samplers/sobol_sampler.py (98%) rename embodichain/lab/sim/{ => motion}/workspace/visualizers/__init__.py (72%) rename embodichain/lab/sim/{ => motion}/workspace/visualizers/axis_visualizer.py (99%) rename embodichain/lab/sim/{ => motion}/workspace/visualizers/base_visualizer.py (99%) rename embodichain/lab/sim/{ => motion}/workspace/visualizers/point_cloud_visualizer.py (98%) rename embodichain/lab/sim/{ => motion}/workspace/visualizers/sphere_visualizer.py (98%) rename embodichain/lab/sim/{ => motion}/workspace/visualizers/visualizer_factory.py (95%) rename embodichain/lab/sim/{ => motion}/workspace/visualizers/voxel_visualizer.py (98%) create mode 100644 embodichain/lab/trajectory_generation/__init__.py create mode 100644 embodichain/lab/trajectory_generation/execution.py create mode 100644 embodichain/lab/trajectory_generation/initial_state.py create mode 100644 embodichain/lab/trajectory_generation/integrations/__init__.py create mode 100644 embodichain/lab/trajectory_generation/integrations/_collision.py create mode 100644 embodichain/lab/trajectory_generation/integrations/atomic.py create mode 100644 embodichain/lab/trajectory_generation/integrations/contact.py create mode 100644 embodichain/lab/trajectory_generation/integrations/planning.py create mode 100644 embodichain/lab/trajectory_generation/integrations/sim.py create mode 100644 embodichain/lab/trajectory_generation/runner.py create mode 100644 embodichain/lab/trajectory_generation/sinks.py rename examples/sim/{ => motion}/planners/curobo_planner.py (98%) rename examples/sim/{ => motion}/planners/neural_planner.py (97%) rename examples/sim/{ => motion}/solvers/differential_solver.py (100%) rename examples/sim/{ => motion}/solvers/neural_ik_solver.py (99%) rename examples/sim/{ => motion}/solvers/opw_solver.py (100%) rename examples/sim/{ => motion}/solvers/pink_solver.py (100%) rename examples/sim/{ => motion}/solvers/pinocchio_solver.py (100%) rename examples/sim/{ => motion}/solvers/pytorch_solver.py (100%) rename examples/sim/{ => motion}/solvers/srs_solver.py (100%) create mode 100644 examples/sim/motion/trajectory_generation/README.md create mode 100644 examples/sim/motion/trajectory_generation/cube_grasp_parallel.py create mode 100644 examples/sim/motion/trajectory_generation/cube_pickup_collection.py create mode 100644 examples/sim/motion/trajectory_generation/free_motion.py rename examples/sim/{ => motion}/workspace/analyze_cartesian_workspace.py (97%) rename examples/sim/{ => motion}/workspace/analyze_joint_workspace.py (97%) rename examples/sim/{ => motion}/workspace/analyze_plane_workspace.py (98%) create mode 100644 tests/gym/envs/test_fixed_scene_preparation.py create mode 100644 tests/lab/trajectory_generation/__init__.py create mode 100644 tests/lab/trajectory_generation/test_atomic_source.py create mode 100644 tests/lab/trajectory_generation/test_contact.py create mode 100644 tests/lab/trajectory_generation/test_cube_grasp_parallel.py create mode 100644 tests/lab/trajectory_generation/test_episode_sinks.py create mode 100644 tests/lab/trajectory_generation/test_execution.py create mode 100644 tests/lab/trajectory_generation/test_initial_state_host.py create mode 100644 tests/lab/trajectory_generation/test_pickup_collection.py create mode 100644 tests/lab/trajectory_generation/test_planning.py create mode 100644 tests/lab/trajectory_generation/test_runner.py create mode 100644 tests/lab/trajectory_generation/test_runner_real.py create mode 100644 tests/lab/trajectory_generation/test_session_sink_integration.py create mode 100644 tests/lab/trajectory_generation/test_sim_initial_state.py rename tests/sim/{ => motion}/planners/test_base_planner.py (96%) rename tests/sim/{ => motion}/planners/test_curobo_integration.py (98%) rename tests/sim/{ => motion}/planners/test_curobo_planner.py (86%) rename tests/sim/{ => motion}/planners/test_motion_generator.py (98%) rename tests/sim/{ => motion}/planners/test_motion_generator_batched.py (98%) rename tests/sim/{ => motion}/planners/test_neural_batched.py (93%) rename tests/sim/{ => motion}/planners/test_neural_planner.py (98%) rename tests/sim/{ => motion}/planners/test_plan_state_batched.py (92%) rename tests/sim/{ => motion}/planners/test_toppra_batched.py (91%) rename tests/sim/{ => motion}/planners/test_toppra_planner.py (96%) rename tests/sim/{ => motion}/solvers/test_base_solver.py (96%) rename tests/sim/{ => motion}/solvers/test_differential_solver.py (100%) rename tests/sim/{ => motion}/solvers/test_neural_ik_solver.py (97%) rename tests/sim/{ => motion}/solvers/test_opw_solver.py (100%) rename tests/sim/{ => motion}/solvers/test_pink_solver.py (99%) rename tests/sim/{ => motion}/solvers/test_pinocchio_solver.py (100%) rename tests/sim/{ => motion}/solvers/test_pytorch_solver.py (100%) rename tests/sim/{ => motion}/solvers/test_srs_solver.py (99%) rename tests/sim/{ => motion}/solvers/test_ur_solver.py (99%) create mode 100644 tests/sim/motion/test_motion_imports.py create mode 100644 tests/sim/motion/trajectory_augmentation/__init__.py create mode 100644 tests/sim/motion/trajectory_augmentation/test_cfg.py create mode 100644 tests/sim/motion/trajectory_augmentation/test_contracts.py create mode 100644 tests/sim/motion/trajectory_augmentation/test_coverage.py create mode 100644 tests/sim/motion/trajectory_augmentation/test_imports.py create mode 100644 tests/sim/motion/trajectory_augmentation/test_operators.py create mode 100644 tests/sim/motion/trajectory_augmentation/test_session.py rename tests/sim/{ => motion}/workspace/__init__.py (100%) rename tests/sim/{ => motion}/workspace/test_analyzer.py (97%) rename tests/sim/{ => motion}/workspace/test_cache.py (97%) rename tests/sim/{ => motion}/workspace/test_runtime.py (98%) diff --git a/.agents/skills/add-solver/SKILL.md b/.agents/skills/add-solver/SKILL.md index 115f46f36..d6d2a1051 100644 --- a/.agents/skills/add-solver/SKILL.md +++ b/.agents/skills/add-solver/SKILL.md @@ -15,7 +15,7 @@ gold standard for structure and style. - User asks to "add a solver", "add an IK solver", "add a kinematic solver" - A new robot family needs a closed-form / numerical / Warp-kernel IK backend - The request names a solver that does not yet exist under - `embodichain/lab/sim/solvers/` + `embodichain/lab/sim/motion/solvers/` ## The Four Artifacts @@ -24,18 +24,22 @@ the conventions below. | # | Artifact | Path | |---|----------|------| -| 1 | Solver module | `embodichain/lab/sim/solvers/_solver.py` | +| 1 | Solver module | `embodichain/lab/sim/motion/solvers/_solver.py` | | 2 | (GPU/Warp only) Warp kernel | `embodichain/utils/warp/kinematics/_solver.py` | -| 3 | Sphinx docs page | `docs/source/overview/sim/solvers/_solver.md` | -| 4 | Unit test | `tests/sim/solvers/test__solver.py` | +| 3 | Sphinx docs page | `docs/source/overview/sim/motion/solvers/_solver.md` | +| 4 | Unit test | `tests/sim/motion/solvers/test__solver.py` | | 5 | Benchmark entry | extend `scripts/benchmark/robotics/kinematic_solver/run_benchmark.py` | Plus two registration edits: - Export the new `Cfg` + `Solver` classes from - `embodichain/lab/sim/solvers/__init__.py`. + `embodichain/lab/sim/motion/solvers/__init__.py`. - Add the docs page to the toctree in - `docs/source/overview/sim/solvers/index.rst`. + `docs/source/overview/sim/motion/solvers/index.rst`. + +Keep solver exports in the solver subpackage. The `motion` parent resolves +subpackages lazily; do not add eager planner or workspace analyzer imports to +the Robot initialization path. ## Steps @@ -53,7 +57,7 @@ Ask the user (only what is not already stated): ### 2. Write the Solver Module -File: `embodichain/lab/sim/solvers/_solver.py` +File: `embodichain/lab/sim/motion/solvers/_solver.py` Required pieces (mirror `ur_solver.py`): @@ -90,7 +94,7 @@ from __future__ import annotations import torch import numpy as np from embodichain.utils import configclass -from embodichain.lab.sim.solvers import SolverCfg, BaseSolver +from embodichain.lab.sim.motion.solvers import SolverCfg, BaseSolver from embodichain.data import get_data_path @@ -172,7 +176,7 @@ For pure-PyTorch / numerical solvers, skip this step entirely and implement ### 4. Register in `__init__.py` Add the import + keep `__all__` (if present) consistent in -`embodichain/lab/sim/solvers/__init__.py`: +`embodichain/lab/sim/motion/solvers/__init__.py`: ```python from .foo_solver import FooSolverCfg, FooSolver @@ -180,7 +184,7 @@ from .foo_solver import FooSolverCfg, FooSolver ### 5. Write the Docs Page -File: `docs/source/overview/sim/solvers/_solver.md` — mirror the structure +File: `docs/source/overview/sim/motion/solvers/_solver.md` — mirror the structure of `ur_solver.md`: 1. `# Solver` — one-paragraph intro (what it solves, why it's fast / what @@ -198,7 +202,7 @@ of `ur_solver.md`: 7. **How It Works** — numbered explanation of the solve pipeline. 8. **References** — markdown links. -Then add the page to the toctree in `docs/source/overview/sim/solvers/index.rst`: +Then add the page to the toctree in `docs/source/overview/sim/motion/solvers/index.rst`: ```rst .. toctree:: @@ -211,7 +215,7 @@ Then add the page to the toctree in `docs/source/overview/sim/solvers/index.rst` ### 6. Write the Unit Test -File: `tests/sim/solvers/test__solver.py` — follow `test_ur_solver.py` +File: `tests/sim/motion/solvers/test__solver.py` — follow `test_ur_solver.py` exactly: - Apache header + `from __future__ import annotations` (after header). @@ -266,16 +270,16 @@ to the `benchmark` skill (`.agents/skills/benchmark/SKILL.md`). ```bash conda activate embodichain -black embodichain/lab/sim/solvers/_solver.py +black embodichain/lab/sim/motion/solvers/_solver.py black embodichain/utils/warp/kinematics/_solver.py # if added -black tests/sim/solvers/test__solver.py +black tests/sim/motion/solvers/test__solver.py black scripts/benchmark/robotics/kinematic_solver/run_benchmark.py ``` Run the unit test (CPU class is enough for a quick check): ```bash -pytest tests/sim/solvers/test__solver.py::TestFooSolver -v +pytest tests/sim/motion/solvers/test__solver.py::TestFooSolver -v ``` Smoke-run the benchmark for the new solver only: @@ -296,7 +300,7 @@ Finally, run the `/pre-commit-check` skill to catch all CI violations locally. - [ ] Solver inherits `BaseSolver`, sets `self.dof`, implements `get_ik` - [ ] `__all__` declared in the solver module - [ ] Google-style docstrings with Sphinx directives (`.. tip::`, `.. attention::`) -- [ ] Cfg + Solver exported from `embodichain/lab/sim/solvers/__init__.py` +- [ ] Cfg + Solver exported from `embodichain/lab/sim/motion/solvers/__init__.py` - [ ] Docs page added to `index.rst` toctree - [ ] Test covers FK↔IK round-trip + unreachable-pose failure, CPU & CUDA classes - [ ] Benchmark entry added with CPU (+CUDA) timing, accuracy, and the 3-table report @@ -308,8 +312,8 @@ Finally, run the `/pre-commit-check` skill to catch all CI violations locally. |---------|-----| | Reimplementing FK / TCP / joint limits in the solver | Reuse `BaseSolver.get_fk`, `set_tcp`, `set_qpos_limits` | | Putting the Warp kernel inside the solver module | Put `@wp.kernel`/`@wp.struct` in `embodichain/utils/warp/kinematics/_solver.py` and import it | -| Not exporting the Cfg/Solver from `__init__.py` | Add the import line so `from embodichain.lab.sim.solvers import FooSolverCfg` works | -| Forgetting the docs toctree entry | Add the `.md` to `docs/source/overview/sim/solvers/index.rst` | +| Not exporting the Cfg/Solver from `__init__.py` | Add the import line so `from embodichain.lab.sim.motion.solvers import FooSolverCfg` works | +| Forgetting the docs toctree entry | Add the `.md` to `docs/source/overview/sim/motion/solvers/index.rst` | | Test only checks happy-path IK | Must verify FK↔IK round-trip equality AND an unreachable-pose returns `False` | | Creating a separate benchmark file | Extend the unified `run_benchmark.py` instead | | Skipping `black` / pre-commit | CI checks every file including tests and benchmarks | @@ -319,13 +323,13 @@ Finally, run the `/pre-commit-check` skill to catch all CI violations locally. | Action | Command / Path | |--------|----------------| -| Reference solver | `embodichain/lab/sim/solvers/ur_solver.py` | +| Reference solver | `embodichain/lab/sim/motion/solvers/ur_solver.py` | | Reference Warp kernel | `embodichain/utils/warp/kinematics/ur_solver.py` | -| Reference docs | `docs/source/overview/sim/solvers/ur_solver.md` | -| Reference test | `tests/sim/solvers/test_ur_solver.py` | +| Reference docs | `docs/source/overview/sim/motion/solvers/ur_solver.md` | +| Reference test | `tests/sim/motion/solvers/test_ur_solver.py` | | Benchmark file | `scripts/benchmark/robotics/kinematic_solver/run_benchmark.py` | | Python env | `conda activate embodichain` | -| Run test (CPU) | `pytest tests/sim/solvers/test__solver.py::TestFooSolver -v` | +| Run test (CPU) | `pytest tests/sim/motion/solvers/test__solver.py::TestFooSolver -v` | | Run benchmark | `python -m scripts.benchmark.robotics.kinematic_solver.run_benchmark -s ` | | Format | `black ` | | Pre-commit | `/pre-commit-check` | diff --git a/.agents/skills/add-test/SKILL.md b/.agents/skills/add-test/SKILL.md index 5536332de..6ff21dc69 100644 --- a/.agents/skills/add-test/SKILL.md +++ b/.agents/skills/add-test/SKILL.md @@ -18,7 +18,7 @@ Write tests following EmbodiChain's conventions and patterns. Tests mirror the source tree under `tests/`: ``` -embodichain/lab/sim/solvers/pytorch_solver.py → tests/sim/solvers/test_pytorch_solver.py +embodichain/lab/sim/motion/solvers/pytorch_solver.py → tests/sim/motion/solvers/test_pytorch_solver.py embodichain/lab/gym/envs/managers/rewards.py → tests/gym/envs/managers/test_reward_functors.py embodichain/toolkits/graspkit/pg_grasp/foo.py → tests/toolkits/test_pg_grasp.py embodichain_tasks/embodichain_tasks/manipulation/push_cube.py → tests/gym/envs/tasks/test_push_cube.py diff --git a/.agents/skills/review-pr/references/review-matrix.md b/.agents/skills/review-pr/references/review-matrix.md index 3627b79ae..468811246 100644 --- a/.agents/skills/review-pr/references/review-matrix.md +++ b/.agents/skills/review-pr/references/review-matrix.md @@ -127,11 +127,14 @@ and the `$add-task-program` read-only deployment inspector. ## Atomic actions, motion planning, and IK -**Paths:** `embodichain/lab/sim/atomic_actions/**`, `planners/**`, -`solvers/**`, and grasp/workspace utilities that feed plans. +**Paths:** `embodichain/lab/sim/atomic_actions/**`, +`embodichain/lab/sim/motion/{planners,solvers,workspace,trajectory_augmentation}/**`, +and grasp utilities that feed plans. Check: +- The `motion` parent and workspace analyzer exports retain lazy loading; + Robot initialization must not eagerly load planners or offline analyzers. - Goal, options, affordance, requirement, binding, plan, command, effect, and evidence types remain coherent across registration, planning, compilation, execution, tracking, and verification. diff --git a/AGENTS.md b/AGENTS.md index 0e5a74cee..3da5820ce 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -70,6 +70,7 @@ EmbodiChain/ │ │ │ ├── runtime/ # Sequential/parallel Semantic Call execution │ │ │ └── integrations/ # Explicit environment and simulation assembly │ │ ├── visualization/ # Browser visualization protocol, runtime, and Viser backend +│ │ ├── trajectory_generation/ # Fixed-scene host, rollout runner, validation adapters and sinks │ │ ├── gym/ # OpenAI Gym-compatible environments │ │ │ ├── envs/ # BaseEnv, EmbodiedEnv, narrow Task Program Gym bridge │ │ │ │ ├── managers/ # Observation, event, reward, record, dataset managers @@ -82,9 +83,11 @@ EmbodiChain/ │ │ │ ├── objects/ # Robot, RigidObject, Articulation, Light, Gizmo, SoftObject │ │ │ ├── sensors/ # Camera, StereoCamera, BaseSensor │ │ │ ├── robots/ # Robot-specific configs and params (dexforce_w1, cobotmagic) -│ │ │ ├── planners/ # Motion planners (TOPPRA, motion generator) -│ │ │ ├── solvers/ # IK solvers (SRS, OPW, pink, pinocchio, pytorch) -│ │ │ └── workspace/ # Reachability analysis and runtime workspace queries +│ │ │ └── motion/ # Robot motion capabilities; subpackages load lazily +│ │ │ ├── solvers/ # IK/FK solvers (SRS, OPW, pink, pinocchio, pytorch) +│ │ │ ├── planners/ # Motion planners (TOPPRA, motion generator) +│ │ │ ├── workspace/ # Reachability analysis and runtime workspace queries +│ │ │ └── trajectory_augmentation/ # Trajectory candidates, operators, coverage, bookkeeping │ │ ├── devices/ # Real-device controllers │ │ └── scripts/ # Environment, preview, and analysis entry points │ ├── toolkits/ # Standalone tools @@ -106,6 +109,15 @@ EmbodiChain/ └── VERSION # Package version file ``` +Import motion APIs from `embodichain.lab.sim.motion.{solvers,planners,workspace,trajectory_augmentation}`. +The `motion` parent resolves subpackages lazily to preserve the +`Robot → solvers` and `planners → SimulationManager` initialization boundary; +workspace analyzer and visualization exports also remain lazy. Trajectory +augmentation algorithms do not directly depend on Gym; execution, reset, and +dataset persistence belong to host integrations. Importing augmentation follows +the normal `lab/sim` initialization lifecycle. Low-level Warp kernels remain +under `utils/warp/kinematics/`. + Official tasks use a task-first layout: - Import-registered Python entry point: `embodichain_tasks/embodichain_tasks//.py` diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 8a6464c77..e66615b56 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -13,6 +13,9 @@ topics: - sim system - lab sim - simulation manager + - motion module + - sim motion + - 运动能力模块 - 仿真系统 - 仿真核心 - 仿真管理器 @@ -52,11 +55,15 @@ topics: - embodichain/lab/sim/objects/__init__.py - embodichain/lab/sim/objects/articulation.py - embodichain/lab/sim/sensors/__init__.py - - embodichain/lab/sim/solvers/__init__.py - - embodichain/lab/sim/planners/__init__.py - - embodichain/lab/sim/workspace/__init__.py + - embodichain/lab/sim/motion/__init__.py + - embodichain/lab/sim/motion/solvers/__init__.py + - embodichain/lab/sim/motion/planners/__init__.py + - embodichain/lab/sim/motion/workspace/__init__.py + - embodichain/lab/sim/motion/trajectory_augmentation/__init__.py - embodichain/lab/gym/envs/base_env.py - embodichain/lab/gym/envs/embodied_env.py + - embodichain/lab/trajectory_generation/initial_state.py + - embodichain/lab/trajectory_generation/integrations/sim.py related_topics: - env-framework - robot-system @@ -86,6 +93,9 @@ topics: - 任务环境 - env profiler - env profiling + - generation lease + - controlled episode preparation + - 受控初态准备 keywords: - env - environment @@ -93,6 +103,10 @@ topics: - embodied_env - EmbodiedEnv - BaseEnv + - acquire_generation_lease + - prepare_generation_episode + - generation_epoch + - FixedSceneHost - register_env - list-task - task listing @@ -124,6 +138,7 @@ topics: - embodichain/lab/gym/envs/demo.py - embodichain/lab/gym/envs/embodied_env.py - embodichain/lab/gym/envs/task_program/bridge.py + - embodichain/lab/trajectory_generation/initial_state.py - embodichain/learning/rl/env.py - embodichain_tasks/embodichain_tasks/ - embodichain/lab/gym/utils/ @@ -222,13 +237,16 @@ topics: paths: - topics/ik-solvers/ik-solvers.md source_of_truth: - - embodichain/lab/sim/solvers/base_solver.py - - embodichain/lab/sim/solvers/srs_solver.py - - embodichain/lab/sim/solvers/opw_solver.py - - embodichain/lab/sim/solvers/pink_solver.py - - embodichain/lab/sim/solvers/pinocchio_solver.py - - embodichain/lab/sim/solvers/pytorch_solver.py - - embodichain/lab/sim/solvers/differential_solver.py + - embodichain/lab/sim/cfg.py + - embodichain/lab/sim/motion/__init__.py + - embodichain/lab/sim/motion/solvers/__init__.py + - embodichain/lab/sim/motion/solvers/base_solver.py + - embodichain/lab/sim/motion/solvers/srs_solver.py + - embodichain/lab/sim/motion/solvers/opw_solver.py + - embodichain/lab/sim/motion/solvers/pink_solver.py + - embodichain/lab/sim/motion/solvers/pinocchio_solver.py + - embodichain/lab/sim/motion/solvers/pytorch_solver.py + - embodichain/lab/sim/motion/solvers/differential_solver.py related_topics: - simulation-system - robot-system @@ -246,12 +264,20 @@ topics: - cobotmagic - control parts - drive properties + - robot workspace + - workspace analysis + - 机器人工作空间 - 机器人系统 - 机器人配置 keywords: - robot - RobotCfg - Robot + - workspace + - RobotWorkspaceCfg + - RobotWorkspace + - WorkspaceAnalyzer + - reachability - control - drive - joint @@ -267,6 +293,10 @@ topics: - embodichain/lab/sim/objects/robot.py - embodichain/lab/sim/robots/ - embodichain/lab/sim/cfg.py + - embodichain/lab/sim/motion/workspace/__init__.py + - embodichain/lab/sim/motion/workspace/cfg.py + - embodichain/lab/sim/motion/workspace/runtime.py + - embodichain/lab/sim/motion/workspace/analyzer.py related_topics: - simulation-system - ik-solvers @@ -388,6 +418,9 @@ topics: - toppra - motion generator - trajectory planning + - trajectory augmentation + - fixed scene trajectory augmentation + - 轨迹扩增 - curobo collision world - dynamic collision integration - 运动生成 @@ -397,6 +430,14 @@ topics: - planner - planning - trajectory + - trajectory_augmentation + - GenerationSession + - GenerationRunner + - FixedSceneHost + - QposRolloutExecutor + - trajectory_generation + - TrajectoryAugmentationCfg + - CandidateTrajectoryBatch - toppra - motion generator - resolve_plan_options @@ -424,13 +465,31 @@ topics: paths: - topics/motion-planning/motion-planning.md source_of_truth: - - embodichain/lab/sim/planners/base_planner.py - - embodichain/lab/sim/planners/toppra_planner.py - - embodichain/lab/sim/planners/curobo/curobo_planner.py - - embodichain/lab/sim/planners/curobo/curobo_yaml.py - - embodichain/lab/sim/planners/motion_generator.py - - embodichain/lab/sim/planners/neural_planner.py - - embodichain/lab/sim/planners/utils.py + - embodichain/lab/sim/motion/__init__.py + - embodichain/lab/sim/motion/planners/base_planner.py + - embodichain/lab/sim/motion/planners/toppra_planner.py + - embodichain/lab/sim/motion/planners/curobo/curobo_planner.py + - embodichain/lab/sim/motion/planners/curobo/curobo_yaml.py + - embodichain/lab/sim/motion/planners/motion_generator.py + - embodichain/lab/sim/motion/planners/neural_planner.py + - embodichain/lab/sim/motion/planners/utils.py + - embodichain/lab/sim/motion/trajectory_augmentation/__init__.py + - embodichain/lab/sim/motion/trajectory_augmentation/contracts.py + - embodichain/lab/sim/motion/trajectory_augmentation/cfg.py + - embodichain/lab/sim/motion/trajectory_augmentation/operators.py + - embodichain/lab/sim/motion/trajectory_augmentation/coverage.py + - embodichain/lab/sim/motion/trajectory_augmentation/session.py + - embodichain/lab/trajectory_generation/runner.py + - embodichain/lab/trajectory_generation/execution.py + - embodichain/lab/trajectory_generation/initial_state.py + - embodichain/lab/trajectory_generation/integrations/sim.py + - embodichain/lab/trajectory_generation/integrations/planning.py + - embodichain/lab/trajectory_generation/sinks.py + - examples/sim/motion/trajectory_generation/free_motion.py + - examples/sim/motion/trajectory_generation/cube_grasp_parallel.py + - examples/sim/motion/trajectory_generation/cube_pickup_collection.py + - embodichain/lab/trajectory_generation/integrations/atomic.py + - embodichain/lab/trajectory_generation/integrations/contact.py - embodichain/lab/task_program/semantics/scene.py - embodichain/toolkits/graspkit/pose_generator.py - embodichain/lab/task_program/integrations/configured.py diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 1e2f623e2..521a41c3d 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -137,6 +137,16 @@ implementation `_plan(request, context)` hook. - `ExecutionRunner` drives that session against observation, command, and clock ports without blocking `step()`. +For selected trajectory candidates, `engine.start(..., +initial_plan_provider=provider)` materializes the first invocation's initial +`ActionPlan` from the newly resolved request and current `PlanningContext`. +The framework still binds collision options, authorizes command destinations, +checks plan identities and scene revisions, and installs tracking and phase +gates. The provider is consumed only during session construction; subsequent +invocations and recovery use the registered skill planner. Rebuild the plan, +binding-dependent effects, and session after reset instead of retaining runtime +state from a previous rollout. + ## PlanningContext invariants `PlanningContext` carries robot observation, scene snapshot, symbolic @@ -180,6 +190,11 @@ Held-object guards and phase-effect gates are observational: Pick gates attachment before lift. Place gates detachment before retract. HandOver owns independent source/destination transfer boundaries. +PickUp's ragged grasp sampling uses an explicit candidate mask. Empty rows and +padding carry a safe current FK pose and cannot win selection; an entirely +empty batch fails without candidate IK. Failed or non-finite IK outputs retain +the preceding valid seed before later pickup stages are screened. + ## Row-local state Vector environments share a synchronized call and command cursor, but success, @@ -271,3 +286,18 @@ python docs/scripts/check_api_docs.py For public API changes also run the docs checker tests and Sphinx dummy build. For simulator adapters add an environment-level test that exercises normal `env.step()` consumption and safe cancellation. + +## Offline PickUp export for fixed-scene collection + +`lab/trajectory_generation/integrations/atomic.py::export_pickup_templates` +exports a successful MoveEndEffector → PickUp compilation into protected phase +qpos templates. It retains approach/close/lift boundaries, expands passive +mimic geometry and appends real hold commands. Only transit permits residuals. +`cube_pickup_collection.py` uses this source with full-state/contact validation +and confirmed LeRobot persistence across repeated full-batch restoration. + +This is an offline atomic source: the qpos executor owns physical validation; +no projected `HeldObjectState` is committed as observed evidence. It does not +consume `initial_plan_provider` or run AtomicActionRuntime tracking/recovery. +The runtime adapter and contact-aware Gym source/host matrix remain separate +acceptance work. Keep the existing runtime and task-state contracts intact. diff --git a/agent_context/topics/env-framework/env-framework.md b/agent_context/topics/env-framework/env-framework.md index d33f85f16..78c66aef9 100644 --- a/agent_context/topics/env-framework/env-framework.md +++ b/agent_context/topics/env-framework/env-framework.md @@ -94,6 +94,68 @@ gym.Env - A structured controller `TensorDict` may carry auxiliary fields such as `ik_success`, but it must contain at least one supported control key. +### Explicit demonstration candidates (`demo.py`) + +- `execute_demo_episode(env, segments=...)` accepts one `DemoSegment` or a lazy + iterable of segments. Explicit segments bypass both task planning factories; + combining them with planning keyword arguments raises `ValueError` before + recording begins. An empty iterable executes no candidate and does not fall + back to task planning. +- `segments=None` preserves `resolve_demo_segments()` and its legacy + `create_demo_action_list()` fallback. Both paths validate segment types lazily + and fill missing instructions from dataset metadata without changing the + supplied segment. +- Explicit candidates use the same normal action processing, per-row masks, + validators, cancellation/abort handshake, and recording lifecycle as task + plans. Auto-reset remains suspended during execution; the caller owns the + later commit or discard boundary. Supplying a candidate does not restore its + initial state. +- Focused coverage: `tests/gym/envs/test_demo.py` and the Task Program bridge + and completion tests under `tests/gym/envs/task_program/`. + +### Fixed-scene generation preparation (`base_env.py`, `embodied_env.py`) + +- `BaseEnv.acquire_generation_lease(owner)` reserves the complete environment + batch by owner identity. Same-owner acquisition is idempotent; other owners + are rejected. The generation flag disables auto-reset independently of demo + and replay flags. Ordinary `reset()` fails before seeding or state mutation + while the lease is held. +- `EmbodiedEnv.prepare_generation_episode(owner, *, prepare, restore, settle, + verify)` is the full-batch discard-and-prepare boundary. Freeze prior episode + evidence first: this method clears camera buffers, expert/trajectory counts, + demo annotations, success state, and the active Task Program bridge without + saving pending episodes. +- Preparation runs deterministic task/controller initialization, physical + restoration, settling, standard manager reset, and nonempty all-passed + `ValidationResult` verification + through trusted callbacks. It does not run startup/reset/interval events or + rewind environment/event RNGs. The host profile owns required initialization + and certifies any interval events that remain active during later rollouts. +- Observation/history, reward, and dataset managers reset before verification; + `get_obs()` and `get_info()` then refresh the initial observation/task state, + and recording seeds from that settled state. Preparation does not use Gym + `step()` or write a training transition. Failure leaves stepping disabled. +- `generation_epoch` advances on lease acquisition/release, every preparation + attempt, and normal resets. Failed preparation invalidates previous epochs. + `FixedSceneHost` in `embodichain/lab/trajectory_generation/initial_state.py` + matches candidate bindings against this epoch and verifies physical/task + initial state before publishing a `PreparedBatch` or Gym first frame. +- `BaseEnv.observe_generation_commands(owner, callback)` observes a successfully + submitted `_step_action` command before physics. The observer requires the + same generation lease and does not issue a second command or step. + `execute_demo_episode` exposes `step_observer(result, active_before_step)` and + `row_step_limits` so the qpos executor records real returned observations and + stops each short row's recording before batch padding/holds. +- `FixedSceneHost.initial_observation(binding)` copies the prepared Gym first + frame without a second `get_obs()` or history update. +- `release_generation_lease(owner)` restores normal reset behavior without + saving, resetting, or modifying demo/replay flags. Callers serialize access; + the host remains the sole simulation stepper while the lease is active. +- Focused coverage: `tests/gym/envs/test_fixed_scene_preparation.py`, including + real Gym lifecycle methods wired to `FixedSceneHost` through fake physical + ports; physical adapters have separate tests under + `tests/lab/trajectory_generation/`. + ### Task Program completion (`embodied_env.py`, `task_program/bridge.py`) - `EmbodiedEnvCfg.task_program` remains opt-in. A registered task may attach diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index 2de569ecc..910e71d2e 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -9,11 +9,11 @@ | File | Role | |---|---| -| `embodichain/lab/sim/solvers/__init__.py` | Public re-exports for all solver classes and configs | -| `embodichain/lab/sim/solvers/base_solver.py` | `BaseSolver` ABC + `SolverCfg` base config | +| `embodichain/lab/sim/motion/solvers/__init__.py` | Public re-exports for all solver classes and configs | +| `embodichain/lab/sim/motion/solvers/base_solver.py` | `BaseSolver` ABC + `SolverCfg` base config | | `embodichain/lab/sim/cfg.py` | `RobotCfg.solver_cfg` — where solver config is wired into a robot | -| `embodichain/lab/sim/solvers/qpos_seed_sampler.py` | `QposSeedSampler` — random joint-seed generation | -| `embodichain/lab/sim/solvers/null_space_posture_task.py` | `NullSpacePostureTask` — Pink null-space posture objective | +| `embodichain/lab/sim/motion/solvers/qpos_seed_sampler.py` | `QposSeedSampler` — random joint-seed generation | +| `embodichain/lab/sim/motion/solvers/null_space_posture_task.py` | `NullSpacePostureTask` — Pink null-space posture objective | | `embodichain/lab/sim/utility/solver_utils.py` | Helpers: `create_pk_serial_chain`, `build_reduced_pinocchio_robot`, `validate_iteration_params`, `compute_pinocchio_fk` | --- @@ -26,6 +26,13 @@ and joint-limit management. A `SolverCfg` subclass is instantiated inside `RobotCfg` and its `init_solver()` factory method produces the concrete `BaseSolver` instance at runtime. +Import solver classes and configs from `embodichain.lab.sim.motion.solvers`. +`RobotCfg.from_dict()` resolves configured `class_type` names against +that public module. The `motion` parent loads subpackages lazily; adding solver +exports must not eagerly import planners or workspace analyzers into the Robot +initialization path. Warp kernels remain under `embodichain/utils/warp/kinematics/`. +Focused solver tests live under `tests/sim/motion/solvers/`. + All solvers use a `pytorch_kinematics` serial chain (`pk_serial_chain`) for FK and Jacobian computation. `torch.compile` is applied to the FK path for performance. diff --git a/agent_context/topics/motion-planning/motion-planning.md b/agent_context/topics/motion-planning/motion-planning.md index e6028a532..27d052dd3 100644 --- a/agent_context/topics/motion-planning/motion-planning.md +++ b/agent_context/topics/motion-planning/motion-planning.md @@ -4,17 +4,26 @@ | What | Path | |---|---| -| Planner registry | `embodichain/lab/sim/planners/__init__.py` | -| Base planner class & config | `embodichain/lab/sim/planners/base_planner.py` → `BasePlanner`, `BasePlannerCfg`, `CollisionWorldInfo`, `PlanOptions`, `validate_plan_options` | -| TOPPRA planner | `embodichain/lab/sim/planners/toppra_planner.py` → `ToppraPlanner`, `ToppraPlannerCfg`, `ToppraPlanOptions` | -| Neural planner | `embodichain/lab/sim/planners/neural_planner.py` → `NeuralPlanner`, `NeuralPlannerCfg`, `NeuralPlanOptions` | -| cuRobo planner | `embodichain/lab/sim/planners/curobo/curobo_planner.py` → `CuroboPlanner`, `CuroboPlannerCfg`, `CuroboWorldCfg`, `CuroboPlanOptions` | +| Planner registry | `embodichain/lab/sim/motion/planners/__init__.py` | +| Base planner class & config | `embodichain/lab/sim/motion/planners/base_planner.py` → `BasePlanner`, `BasePlannerCfg`, `CollisionWorldInfo`, `PlanOptions`, `validate_plan_options` | +| TOPPRA planner | `embodichain/lab/sim/motion/planners/toppra_planner.py` → `ToppraPlanner`, `ToppraPlannerCfg`, `ToppraPlanOptions` | +| Neural planner | `embodichain/lab/sim/motion/planners/neural_planner.py` → `NeuralPlanner`, `NeuralPlannerCfg`, `NeuralPlanOptions` | +| cuRobo planner | `embodichain/lab/sim/motion/planners/curobo/curobo_planner.py` → `CuroboPlanner`, `CuroboPlannerCfg`, `CuroboWorldCfg`, `CuroboPlanOptions` | | Planner assets | `embodichain/data/assets/planner_assets.py` → `download_neural_planner_checkpoint()` | -| Motion generator | `embodichain/lab/sim/planners/motion_generator.py` → `MotionGenerator`, `MotionGenCfg`, `MotionGenOptions` | -| Planner utilities & data types | `embodichain/lab/sim/planners/utils.py` → `PlanState`, `PlanResult`, `MoveType`, `MovePart`, `TrajectorySampleMethod`, `interpolate_xpos_batched` | +| Motion generator | `embodichain/lab/sim/motion/planners/motion_generator.py` → `MotionGenerator`, `MotionGenCfg`, `MotionGenOptions` | +| Planner utilities & data types | `embodichain/lab/sim/motion/planners/utils.py` → `PlanState`, `PlanResult`, `MoveType`, `MovePart`, `TrajectorySampleMethod`, `interpolate_xpos_batched` | +| Trajectory augmentation | `embodichain/lab/sim/motion/trajectory_augmentation/` → contracts, configs, operators, coverage, `GenerationSession` | ## Overview +Planners, solvers, workspace analysis, and trajectory augmentation are sibling +packages under `embodichain.lab.sim.motion`. Import planner APIs from +`embodichain.lab.sim.motion.planners`; the parent namespace resolves subpackages +lazily and must not eagerly load planners during Robot initialization. Atomic +Actions consume these motion capabilities from `sim/atomic_actions/`. +Focused tests and examples live under `tests/sim/motion/` and +`examples/sim/motion/`. + The planning stack has two layers: 1. **BasePlanner** — low-level trajectory planner that takes a list of `PlanState` waypoints and produces a `PlanResult` with joint trajectories. 2. **MotionGenerator** — the single stateful planning facade that composes a @@ -25,6 +34,117 @@ All planners resolve their robot at init via `SimulationManager.get_instance().g The entire stack is **env-batched** (`B = num_envs`). `PlanState` / `PlanResult` tensors carry a leading `B` dimension; `BasePlanner.plan()` and `MotionGenerator.generate()` operate on `B` environments in one call. +## Trajectory Augmentation Boundary + +`motion/trajectory_augmentation/` owns immutable trajectory/candidate contracts, +strict configuration decoding, phase-authorized joint residuals and retiming, +sampled motion-limit checks, geometric/timing coverage, and bounded generation +session bookkeeping. Candidates are logical rows; their identities and local +random seeds do not derive from physical environment slots. + +`GenerationSession` accounts for ready candidates, rollout and pending-write +budgets, accepted-episode coverage reservations, and idempotent commit receipts. +It never steps or resets an environment or writes a dataset. Execution, +initial-state restoration, physical/task validation, and durable sinks belong +to host integrations. Keep algorithm modules free of direct Gym imports while +recognizing that the public package follows normal `lab/sim` initialization. +Motion-limit validation alone does not establish collision freedom or task +success. + +`rotate_grasp_about_object_axis` rotates a reference TCP pose about a fixed +object-local axis through the object origin. The caller chooses geometry-valid +angles and replans the resulting pose candidates; the operator neither moves +the object nor certifies the grasp. A cube's quarter-turn symmetry is used by +`examples/sim/motion/trajectory_generation/cube_grasp_parallel.py` to compare +0/90-degree grasp orientations crossed with reference/residual transit paths. +Four physics rows execute together using existing MoveEndEffector/PickUp plans; +only the pre-grasp transit allows `joint_residual`. The example saves a tiled +camera MP4 with measured TCP trails, `rollout.npz`, and `report.json`. Physical +acceptance checks sustained lift and TCP-relative position stability. It uses +IK interpolation and does not claim contact-aware collision validation or +LeRobot expert qualification; the separate `cube_pickup_collection.py` performs +that bounded collection workflow. + +`lab/trajectory_generation/runner.py` owns the synchronous qpos job for +handwritten free motion and offline atomic PickUp exports. It resolves supplied host/planner/executor/sink policy identities, +reserves bounded rollout capacity through `GenerationSession`, and restores the +full batch only when ready candidates need another round. The concrete +`QposRolloutExecutor` uses normal Gym demo/controller ports or one pure-sim +step owner, captures actual submitted targets and T+1 observations, and checks +the manager's actual control-time increments. Measured qpos is rechecked for +collision, dynamics, reference-relative length/duration, and task success. +Only `LeRobotEpisodeSink` receipts confirmed after sealed readback update +committed counts and coverage; the runner closes its host/sink and writes +`generation_report.json` on completion or failure. + +`integrations/planning.py` exposes `EEFPath` and `EnvRowMotionPlanner`: logical +candidate count C is scheduled over real physical B rows, without changing the +backend's batch size. It supports fully annotated free paths, exact qpos checks, +EEF IK/FK conversion and explicitly supplied solved branches. Contact, held +objects, changing locked joints, and unavailable backend checks cannot pass. +Joint-path collision checks use bounded sample densification, not a continuous +collision guarantee. The cuRobo lock model cache includes actual locked-joint +initial values and rejects runtime configuration drift until explicitly closed. + +Generation currently requires every physical rigid UID in both the collision +world and its pose-update IDs, so each initial-state snapshot supplies the +correct per-row geometry pose. Different relative layouts require a per-env +world. The executor certifies fixed root/rigid poses at observation boundaries; +implicit floor geometry and other fixed conditions remain the trusted profile's +responsibility. This free-motion integration does not qualify atomic PickUp, +attachment geometry, or the planned four source/host combinations. + +Focused integration tests are in `tests/lab/trajectory_generation/`; use the +CPU tests for orchestration and resource limits, and opt into the serialized +real cuRobo/physics smoke tests when changing those adapters. + +`examples/sim/motion/trajectory_generation/free_motion.py` runs the complete +normal-gravity pure-arm UR5 pipeline and emits a LeRobot shard plus report. +`--record-video --duration 5 --joint-displacement 0.4` also streams a 640x480, +20 fps offscreen camera preview to `preview.mp4`. Rendering observes the actual +rollout without stepping physics; video includes both initial and terminal +observations and stays separate from the numeric LeRobot schema. +`--robot panda` is a real rejection case: motion-induced hand/mimic displacement +exceeds the locked-joint model tolerance, so actual collision verification is +unavailable and no expert episode is committed. These two real cases are in +`tests/lab/trajectory_generation/test_runner_real.py`; they do not replace the +planned PickUp source/host qualification matrix. + +## PickUp Collection Integration + +`integrations/atomic.py::export_pickup_templates` accepts a successful offline +MoveEndEffector → PickUp compilation, preserves its explicit phase boundaries, +expands mimic coordinates, and adds real terminal hold commands. Only transit +allows residuals. It exports qpos references, without committing projected +symbolic effects or executing AtomicActionRuntime recovery/tracking. + +`integrations/contact.py::PickUpMotionValidator` is the shared Runner planner +and executor contact validator. It uses conservative URDF collision hulls, +full-joint FK, cuboid world geometry, an implicit floor, phase/contact-entry +permissions, and planned held-object geometry. Actual collision checks use +measured fingers/object poses. Native contacts are checked after every 5 ms +CPU physics step; forbidden/unknown/cross-row/deep contacts, buffer overflow, +missing finger contact, slipping and dropping cannot qualify. Mounting links +must be fixed to the robot root. Two-hop self-pair exclusions match cuRobo. +The executor exempts only this target and declared mimic joints from its +free-motion immobility rules; all remaining root/world checks remain active. + +`examples/sim/motion/trajectory_generation/cube_pickup_collection.py` runs four +rows, repeatedly restores the full initial state, and collects confirmed +LeRobot shards with optional synchronized video. Its 0/90-degree targets use +independent residual transit paths. Controller stiffness/opening margin are +set before capture; collision and tracking thresholds are not relaxed per row. +The sink flattens numeric matrices in C order, records `observation_shapes`, +and preserves terminal matrices. `commanded_joint_indices` distinguishes active +controller labels from full-joint passive target columns. + +This is CPU-physics / unscaled fixed-base URDF / cuboid-rigid support. The +optional `trajectory-generation` dependencies supply FCL/trimesh/yourdfpy. +It is sampled validation, not continuous collision detection. Contact-aware +Gym, atomic runtime replay, general via points and a registry/CLI remain work. +Focused tests: `test_contact.py`, `test_atomic_source.py`, `test_episode_sinks.py`, +`test_pickup_collection.py` under `tests/lab/trajectory_generation/`. + ## Planner Hierarchy ``` @@ -295,7 +415,7 @@ total duration and emits new explicit arrival intervals. "neural": (NeuralPlanner, NeuralPlannerCfg), } ``` -6. Export from `embodichain/lab/sim/planners/__init__.py`. +6. Export from `embodichain/lab/sim/motion/planners/__init__.py`. ### validate_plan_options decorator diff --git a/agent_context/topics/robot-system/robot-system.md b/agent_context/topics/robot-system/robot-system.md index 3fc36e140..6f17cc2f9 100644 --- a/agent_context/topics/robot-system/robot-system.md +++ b/agent_context/topics/robot-system/robot-system.md @@ -11,6 +11,9 @@ | Robot registry (all robots) | `embodichain/lab/sim/robots/__init__.py` | | DexforceW1 config package | `embodichain/lab/sim/robots/dexforce_w1/` | | CobotMagic config | `embodichain/lab/sim/robots/cobotmagic.py` | +| Solver APIs | `embodichain/lab/sim/motion/solvers/__init__.py` | +| Workspace runtime and config | `embodichain/lab/sim/motion/workspace/runtime.py`, `cfg.py` | +| Workspace offline analysis | `embodichain/lab/sim/motion/workspace/analyzer.py` | | Add-robot tutorial | `docs/source/tutorial/add_robot.rst` | | Add-robot quick-reference | `docs/source/guides/add_robot.rst` | @@ -23,6 +26,23 @@ A `Robot` is instantiated with a `RobotCfg` and a list of DexSim `Articulation` entities. +## Motion and Workspace Integration + +- `Robot` imports solver APIs and runtime workspace types from + `embodichain.lab.sim.motion`. `RobotCfg.from_dict()` resolves solver + `class_type` names against `embodichain.lab.sim.motion.solvers`. +- `RobotCfg.workspace_cfg` maps control-part names to `RobotWorkspaceCfg`. + `Robot.get_workspace(name)` lazily loads the configured cache on first use; + robot construction does not read workspace cache files. +- `Robot.attach_workspace()` accepts a prebuilt `RobotWorkspace` and moves it + to the robot device. `Robot.sample_reachable_pose()` uses the runtime cache + to return bounded batched workspace samples. +- `motion` subpackages load lazily, and workspace analyzer/visualization APIs + retain a separate lazy export boundary. Do not introduce offline analysis + imports into the Robot initialization path. +- Focused workspace coverage lives under `tests/sim/motion/workspace/`; solver + tests live under `tests/sim/motion/solvers/`. + ## RobotCfg Pattern Inheritance chain: diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 6010e0dd1..07573cd50 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -5,6 +5,7 @@ | What | Path | |------|------| | Public simulation package | `embodichain/lab/sim/__init__.py` | +| Motion capability namespace | `embodichain/lab/sim/motion/__init__.py` | | World and scene owner | `embodichain/lab/sim/sim_manager.py` → `SimulationManager` | | Global simulation config | `embodichain/lab/sim/sim_manager.py` → `SimulationManagerCfg` | | Object and physics configs | `embodichain/lab/sim/cfg.py` | @@ -13,7 +14,8 @@ `embodichain.lab.sim` exports the manager, its config, shared material types, `BatchEntity`, and the simulation profiler. Import a specialized -object, sensor, solver, planner, or atomic-action API from its own subpackage. +object, sensor, or atomic-action API from its own subpackage. Solver, planner, +workspace, and trajectory-augmentation APIs live under `embodichain.lab.sim.motion`. ## Ownership @@ -95,17 +97,47 @@ flag later for all or selected environment indices. | Rigid, deformable, articulation, robot, light, constraint, gizmo | `objects/` | `robot-system` for robots | | Camera, stereo camera, contact sensor | `sensors/` | `sensor-system` | | Robot-specific configuration | `robots/` | `robot-system` | -| Inverse kinematics | `solvers/` | `ik-solvers` | -| Trajectory and motion generation | `planners/` | `motion-planning` | +| Inverse kinematics | `motion/solvers/` | `ik-solvers` | +| Trajectory and motion generation | `motion/planners/` | `motion-planning` | +| Trajectory candidates, augmentation, coverage, generation bookkeeping | `motion/trajectory_augmentation/` | `motion-planning` | | Typed action planning and execution | `atomic_actions/` | `atomic-actions` | | Task Program Semantic Calls and robot profiles | `embodichain/lab/task_program/semantics/` | `task-programs` | -| Reachability analysis and runtime workspace queries | `workspace/` | `robot-system` | +| Reachability analysis and runtime workspace queries | `motion/workspace/` | `robot-system` | | Browser scene export and Viser runtime | `embodichain/lab/visualization/` | `sim-visualization` | Use the narrow topic when a request names one of these subsystems. Use `simulation-system` for the overall `lab/sim` architecture, manager lifecycle, scene ownership, or cross-module flow. +`motion/__init__.py` resolves its four public subpackages lazily. Keep it free +of eager planner imports: `Robot` needs solver and runtime workspace types +during initialization, while planners resolve robots through +`SimulationManager`. Workspace analyzer/visualization exports retain their own +lazy boundary. Import concrete APIs from the corresponding `motion` subpackage; +the former `sim.solvers`, `sim.planners`, and `sim.workspace` paths have no +compatibility packages. + +Trajectory augmentation owns value contracts, operators, coverage, and session +bookkeeping. Its algorithms do not import Gym or perform environment execution, +reset, or dataset I/O; those operations belong to host integrations. Its public +import follows the normal `lab/sim` lifecycle and does not promise an isolated +toolkit import. Atomic Actions remain above the motion capabilities. + +`FixedSceneHost` in `embodichain/lab/trajectory_generation/initial_state.py` +reserves an entire simulator batch for capture/restore and publishes only +validated `PreparedBatch` epochs. It compares both the physical adapter's +structure signature and the trusted profile's fixed-condition signature; +acquisition verifies the captured physical state before accepting it. Normal +`SimulationManager.reset_objects_state()` calls fail while this host owns the +batch. Its physical adapter restores state directly, and a Gym-backed host +also uses the controlled generation lease/preparation boundary described in +`env-framework`. + +`SimulationManager.simulation_time` is the read-only sum of successful manual +`update()` physics steps. Generation records elapsed control time from this +counter; preparation setters that update the native world directly are outside +this clock and occur before the rollout timestamp origin is captured. + `Articulation.get_parent_joint_chain(link_name)` is the public topology query for integrations that need link ancestry. It returns immediate-parent-first `ArticulationJointKinematics` values containing copied names, joint type, diff --git a/docs/design/fixed_scene_expert_trajectory_augmentation_design.md b/docs/design/fixed_scene_expert_trajectory_augmentation_design.md new file mode 100644 index 000000000..08db0968b --- /dev/null +++ b/docs/design/fixed_scene_expert_trajectory_augmentation_design.md @@ -0,0 +1,578 @@ +# EmbodiChain：固定场景下的专家轨迹域扩增设计 + +状态：设计与实施边界，更新于 2026-09-05。手写 qpos/free-motion 同步闭环已落地；完整配置、接触/持物与调度方案仍含待实现内容，当前边界见第 11 节及[实施计划](fixed_scene_expert_trajectory_augmentation_implementation_plan.md)。本文没有吞吐对比实测，不宣称某种模式在所有任务上最优。 + +适用范围：原子技能、手写关节轨迹、EEF waypoint 和可信参数化工厂共享扩增能力，由 sim 或 Gym 调用。场景布局与物体初始位姿由宿主提供，不属于扩增变量。 + +推荐架构：**按场景隔离的候选池 + 可选同场景副本池 + 分层按需扩增 + 成本感知调度**。首版保留全批执行边界，通过时长分桶、计算缓存和有界队列提高效率；逐行异步恢复、检查点分叉作为能力成熟后的选项。 + +## 1. 目标、边界与效率指标 + +### 1.1 固定的是每个场景条件，不是所有环境 + +一次扩增的输入为 `(S, G, Q₀)`： + +| 输入 | 固定内容与允许变化 | +|---|---| +| 场景 S | 实体及物体初始位姿、几何、物理参数、机器人安装位姿；不由扩增器采样 | +| 任务 G | 操作对象、目标效果、阶段依赖、接触约束与成功条件；不改变任务语义 | +| 起点集合 Q₀ | 默认使用调用者给定的完整机器人初态；只有部署允许时才扩增关节起点 | + +多环境可以具有不同的 S；同一环境在外层 reset 后也可以接收新的 S。扩增器只计算各条件下的轨迹分布，不决定如何生成布局。任务执行导致物体运动属于正常任务动力学，不是场景随机化。 + +目标覆盖“从允许起点能到达，并能继续完成任务”的状态—动作路径。状态包含关节构型、阶段与接触关系；只证明某个 EEF 点运动学可达不够。[Robot 的工作空间采样](../../embodichain/lab/sim/objects/robot.py) 可作先验，但仍需验证前缀、后缀与整条路径。 + +### 1.2 优化目标 + +先满足任务、碰撞、控制和专家质量门槛，再在预算内增加有效覆盖： + +- 已验收且已持久化的去重轨迹数／墙钟时间。 +- 实际轨迹新增的几何、构型与时间覆盖／墙钟时间。 +- 达到指定覆盖目标的总成本，包含准备、规划、物理、渲染、验收和写盘。 + +规划成功数、种子数和时间变体数不能代替上述指标。无意义绕路、长停顿、近重复不因增加点数获得奖励。比较方案时固定硬件、传感器、任务质量门槛和数据目标，并分别报告冷启动与稳定运行成本。 + +## 2. 分层架构与来源接入 + +### 2.1 所有权与依赖方向 + +| 层 | 职责 | 边界 | +|---|---|---| +| 公共扩增核心 | 模板、候选值对象、因子算子、去重、覆盖描述符与规划端口协议 | 算法直接依赖 tensor/math/configclass，不直接导入 Gym 或仿真 backend；公共导入遵循 lab/sim 初始化 | +| GenerationSession | 各 case 的候选池、预算、随机流、覆盖预留及提交计数 | 只接收快照和结果，不读取活环境，不执行 step/reset | +| 来源与规划适配 | 将原子或手写输入转为模板，调用 IK、MotionGenerator 与路径检查 | 注入服务；原子 binding/effect 不泄漏到公共核心 | +| Runner 与宿主适配 | 执行槽管理、初态准备、命令、步进、观测、效果验证和数据提交 | 每个宿主只有一个物理步进所有者 | + +公共核心位于 `embodichain/lab/sim/motion/trajectory_augmentation/`，包含 `contracts.py`、`cfg.py`、`operators.py`、`coverage.py`、`session.py`，与 `motion/{solvers,planners,workspace}` 共属机器人运动能力域。已实现的 `initial_state.py`、`runner.py`、`execution.py`、`sinks.py` 和 `integrations/{sim,planning}.py` 位于 `embodichain/lab/trajectory_generation/`；Gym 生命周期接入现有 env，原子来源适配仍待实施。端口协议集中在 contracts，避免过早拆分。 + +2026-09-05 根据模块归属 review，将原工具包方案调整为 `lab/sim/motion`：`motion` 父包按需加载四个子包，不在初始化时主动汇总全部求解器、规划器或分析器。当前 [lab 包初始化](../../embodichain/lab/__init__.py) 仍会加载仿真相关子包，因此不再承诺无需仿真依赖即可公共导入；纯算法不直接依赖 Gym，也不拥有 step/reset。sim/Gym 必须复用同一生成会话，不各自实现采样、预算和覆盖逻辑。求解器不反向依赖扩增会话；Gym 生命周期适配不进入 sim 运动核心。 + +### 2.2 原子技能与手写轨迹同等接入 + +| 来源 | 必需输入 | 适配方式 | +|---|---|---| +| 原子技能／Task Program | invocation、阶段约束、affordance、binding 与效果条件 | 导出模板，保留候选到规划完成；只对执行分支构造 ActionPlan | +| 手写 `qpos + dt` | joint names/单位、工具通道、阶段、锁定端点和可变范围 | 扩增声明的自由段，保留未受控关节；未知接触段不改动 | +| 手写 EEF waypoint | 坐标系、工具标定、阶段与路径约束 | 经注入的 IK/规划端口生成关节分支 | +| 可信参数化工厂 | 显式快照、因子值、局部 RNG，以及模板／下一阶段输出 | 概念接口 `build(snapshot, factors, rng) -> TrajectoryTemplate` | + +未标注的 `qpos + dt` 默认仅支持原样执行、记录和审计。单纯 action 数组也不能自动解释为 qpos；velocity、EEF delta 等需声明控制语义及可靠转换，否则不启用几何扩增。 + +现有 `create_demo_action_list()`／`create_demo_segments()` 可适配，但会 step 或依赖未来反馈的旧生成器，应先做受控基准采集，或按实际阶段边界生成下一段；不能在纯规划核心提前耗尽 lazy iterator。配置只引用可信注册的 source ID，不执行任意 YAML 表达式或 dotted import。 + +### 2.3 sim/Gym 执行接入 + +| 来源 × 宿主 | 执行路径 | +|---|---| +| 原子 × sim | 普通 ActionPlan → 现有原子 ExecutionSession/Runner 与仿真端口 | +| 原子 × Gym | Task Program/Gym 命令桥 → 正常 `env.step()` | +| 手写 × sim | 计时控制命令与工具事件 → sim 适配器；无需伪造 AtomicAction 或 Gym | +| 手写 × Gym | 显式候选 DemoSegment → demo 执行器 → 正常 `env.step()` | + +Gym 尚需补充显式 segments 输入或明确消费候选的工厂入口，不能假定旧方法的 kwargs 会生效，也不临时替换环境方法。已是控制器目标的命令走 [ControllerAction](../../embodichain/lab/gym/envs/types.py)，只跳过重复 raw-action pre 处理，不绕过 step 或 post 处理。 + +Gym 的权威周期为 `env.step_dt`；直接 sim 显式提供 `control_dt`，并验证对应整数个 physics step。Gym Runner 不额外调用 sim.update；直接 sim 由适配器完成控制、物理、传感器与记录。规划使用不可变快照；后台状态改变后，受影响计划必须重新检查。 + +## 3. 统一数据契约与身份 + +### 3.1 核心值对象 + +| 契约 | 关键字段与约束 | +|---|---| +| `SceneCase` | 场景内容签名、任务/机器人/标定契约、允许起点；不持有活环境引用 | +| `MotionSnapshot` | case ID、完整机器人状态、根坐标系、实体状态、目标及依赖 revision；输入行到宿主实例的映射由适配器维护 | +| `TrajectoryTemplate` | source ID/revision、template ID、控制表示、参考关节或 EEF 路径、显式时间或待参数化声明、阶段/锚点、允许算子、工具事件和 validator ID | +| `CandidateTrajectoryBatch` | `positions: (C,N,D_full)`、`dt: (C,N)`、`valid_length: (C,)`；每行 case/初态/候选身份、实际 IK 分支、因子值、阶段事件、验证状态及诊断 | +| `ExpertEpisode` | 实际观测、下发动作与表示、时间、机器人/接触测量、阶段、验收证据、候选谱系及提交身份 | + +C 为逻辑候选数，不等于物理环境数；N 可 padding，只有 valid_length 内的样本有效。适配器转换现有 PlanResult、TimedTrajectory 和 ActionPlan；这些 lab 类型不进入核心协议。失败候选可删除大张量，但必须保留轻量 audit。 + +`GraspCandidateBatch` 仍由独立 graspkit 所有：`poses: (B,K,4,4)`、`costs: (B,K)`、`candidate_mask: (B,K)`。兼容辅助方法可包装原 ragged 接口；镜像、upright、标定及语义过滤由来源适配器派生。验证有限值及 SE(3) 合法性；空行填安全 pose 并置 mask=False，不访问不存在的第一个 grasp。 + +### 3.2 三种身份不能混用 + +| 身份 | 含义与生命周期 | +|---|---| +| `scene_case_id / initial_state_id` | 哪个固定场景与允许起点；相同 case 可服务多个副本及多次 episode | +| `candidate_id / geometry_family_id / parent_id` | 哪条扩增方案、几何家族和父分支；重试增加 attempt ID，时间变体保留同一几何家族 | +| `slot_id / runtime_epoch` | 本次在哪个物理实例执行，以及该槽当前的运行时版本；由宿主绑定 | + +`source_row_index` 仅表示输入快照行，不能当作永久 env 身份；来源 env ID 可作为宿主审计信息另存。同一候选转移到副本前要检查场景、起点、控制和坐标映射兼容性；reset 后重建运行时 binding、效果请求和碰撞依赖。 + +离线可返回全部规划合格候选；在线兼容路径仍逐环境按可行性、质量分数及原始候选顺序稳定选择一个 winner。无解行保持起点并标记失败，不挤占其他行。不得将重复 env IDs 或 `(B,K,N,D)` 塞入现有 [TimedTrajectory](../../embodichain/lab/sim/atomic_actions/plans.py)。手写来源没有 grasp 时相应字段为空，不伪造索引。 + +## 4. 扩增维度与分层生成 + +### 4.1 统一因子表 + +| 因子 | 采样变量与粒度 | 必要条件／不变量 | +|---|---|---| +| `contact` | 每个接触分支的 grasp 区域、合法对称变体、夹持宽度 | affordance 与可重新生成接触段的模板；保持任务和物体初态 | +| `ik` | 每个目标/路径的肘部、腕部、冗余构型 | EEF 约束，或明确许可从 qpos 经 FK 重建；关节限位与连续可达 | +| `approach` | 每条几何候选的接近距离、许可方向、撤离方式 | 固定接触锚点、进入方向及合法走廊 | +| `spatial` | 自由段的少量路点、平滑残差、姿态过渡、绕障走向 | 已声明自由段、锁定端点及质量边界；不逐控制帧重采样 | +| `timing` | 每条几何路径的分阶段速度和加减速、时长 | 统一 control_dt，动态限制与阶段连接 | +| `contact_timing` | 每个闭合/释放/稳定事件的许可时长 | 工具通道、事件顺序、接触模式和效果 gate | +| `start_state`，可选 | episode 准备前选择机器人起始构型 | 起点属于部署分布，场景与机器人安装位姿不变 | +| `recovery`，后续 | 从实际偏离状态生成成功后缀 | 真实可达的前缀、后缀生成器与恢复验收 | + +采样域为“任务允许域 ∩ 模板允许域 ∩ job 配置 ∩ 机器人/backend 能力”。缺少启用因子的标注或交集为空时明确报错；关闭因子保持模板值。先固定起点，条件采样接触/接近方式，再联合求解路径与 IK，最后按需求产生时间变体；不展开无界 Cartesian product。 + +接触候选先过质量门槛，再按接触区域、方向等分层选代表，不只按全局 grasp cost top-K。未标定 cost 不是成功概率。一个自由段首版选一种几何算子；多算子组合需明确顺序并重验最终路径。 + +### 4.2 IK、几何与时间的关键约束 + +1. **多解 IK**：解析枚举或多 seed 求解后按关节构型去重；每分支独立传播 pre-grasp → grasp → lift → downstream 的成功 seed。NaN/Inf、FK residual 越界均失败，失败 qpos 不污染后续阶段。 +2. **保持实际分支**：请求携带 `solved_joint_targets`、分支及连续性约束。现有 [EEF 插值路径](../../embodichain/lab/sim/motion/planners/motion_generator.py) 会再次做 IK；只传 pose 可能使多分支重新合并。无法强制保持时按最终结果重新识别、去重。 +3. **低维几何变化**：自由段用少量 via points 或 `q(s)=q_ref(s)+Σ θ_j B_j(s)`，基函数保持端点及所需导数；避免逐帧独立噪声。接触段必须遵守 Cartesian/接触约束，不能只做关节直线插值。 +4. **显式时间参数化**：同一路径改变推进速度，再按 control_dt 生成变长命令；不能同时固定点数、周期又改变总时长。重新计算事件索引，检查速度、加速度及任务所需 jerk/力矩。未许可的工具与等待时长不自动缩放。 + +同 seed、不同 seed、不同 chunk、不同点数都不自动代表不同运动模式。时间变体可增加速度覆盖，不增加同一路径的空间覆盖。 + +### 4.3 多阶段与技能推广 + +分支节点至少保留完整机器人状态、夹爪/接触状态、持物关系、语义任务状态、阶段、父前缀及后缀约束。不能仅因 qpos 接近就合并两个节点。 + +若 T 表示后一个坐标系在前一个坐标系中的位姿，候选 i 的下游 EEF 目标为： + +$$ +T^{(i)}_{\mathrm{object}\to\mathrm{eef}} += T_{\mathrm{world}\to\mathrm{object}}^{-1}T^{(i)}_{\mathrm{world}\to\mathrm{eef}}, +\qquad +T^{(i)}_{\mathrm{world}\to\mathrm{eef,target}} += T_{\mathrm{world}\to\mathrm{object,target}}T^{(i)}_{\mathrm{object}\to\mathrm{eef}}. +$$ + +只对未来可继续、质量合格的节点投入后续预算；物理执行后的续段使用实测状态，不把计划终点当作已实现状态。 + +- PickUp：优先复用候选 IK，延后 winner；闭合、lift 和后续放置继承同一分支。 +- AxisAlign：把对齐偏好变为候选评分，两阶段及持物旋转均保留候选依赖。 +- HandOver：两侧先筛选有限兼容 pair，按 control part/assignment 分组；双臂各自可达不等于联合路径、交接时序可行。 +- Slide、OpenDoor 及手写多段轨迹复用阶段协议；接触候选来自哪个接口不应限制通用扩增。lift_height 等参数只有被定义为可变中间量时才能采样。 + +### 4.4 避障驱动的空间扩增 + +固定场景几何是规划条件，不是扩增变量。空间扩增分为两种模式，避免把“检查后剔除碰撞路径”等同于“主动生成不同绕障路线”： + +| 模式 | 生成方式 | 用途 | +|---|---|---| +| 局部扩增与过滤 | 在参考自由段附近采样受限残差或路点，再检查整段路径 | 低成本增加参考路线附近的覆盖 | +| 避障重规划 | 固定任务锚点、阶段边界及声明的 IK 分支约束，采样许可路点或规划初值，调用具备避障能力的规划端口 | 探索不同的有效绕行路线，不保证不同 seed 得到不同路线 | + +避障重规划是候选生成方式，不仅是碰撞失败后的补救；`ik_interp` 或单纯时间参数化不能替代它。原子与手写来源复用同一规划端口，只改模板授权的自由段，不放松接触锚点、后缀可行性或专家质量门槛。 + +- **有限修复**：碰撞候选可在预算内重规划或换路点；显式限制每候选重规划次数和每 case 规划预算，耗尽后记录失败。启用模式但缺少所需后端能力时明确拒绝，不静默退回插值。 +- **最终验收**：检查平滑、拼接、时间重采样后的实际输出路径，并更新其阶段事件与谱系。碰撞世界绑定和接触/持物规则复用第 5.2、8.1 节；能力不可用不计为通过，不能只检查修复前路径或关键帧。 +- **有效覆盖**：按最终几何路线去重；多个提议修复到同一路径时不重复计空间覆盖。安全间隙与合理路径长度作为质量门槛，不奖励无意义绕路。 + +上述模式选择与重规划预算属于待补配置。当前 `EnvRowMotionPlanner` 主要执行给定 EEF 路径的 IK/FK 转换及 free/no-held 关节路径加密检查，尚不主动搜索替代绕障路线；接触、持物和连续碰撞的实现边界仍见第 11 节。 + +## 5. 高效批量规划与缓存 + +### 5.1 逻辑候选与 GPU 批次分离 + +存储使用紧凑候选表；执行使用少量预热的固定容量桶，例如 16/32/64 行。按机器人/控制部位、阶段结构、碰撞配置和长度分组，桶内 padding 使用安全输入与显式 mask；无效行即使 backend 返回成功也不得进入候选池。 + +每阶段压到最小 C 不一定最快：改变 batch/horizon 可能增加分配、warmup 和 CUDA Graph 成本。优先在阶段边界整理候选，控制桶数量及缓存显存，而不是为每个 C 创建 backend。cuRobo 支持容量预分配和内部 padding，项目已有 [按 batch 等条件缓存 backend](../../embodichain/lab/sim/motion/planners/curobo/curobo_planner.py) 的基础，不应重复建设。[cuRobo 批量规划说明](https://nvlabs.github.io/curobo/latest/api/curobo.batch_motion_planner.html) + +建议由 `generate_candidates()` facade 负责能力检查、分桶与以下映射： + +- 候选 → 来源 snapshot/case → 机器人根坐标与完整起始状态。 +- 候选 → 当前阶段碰撞世界、允许接触对及持物几何。 +- 规划行 → 原 candidate ID、阶段 mask、实际分支和失败原因。 + +[BasePlanner 校验](../../embodichain/lab/sim/motion/planners/base_planner.py) 当前仍将 batch 绑定到 robot.num_instances，且 [机器人 IK](../../embodichain/lab/sim/objects/robot.py) 存在按环境读取根坐标的路径。不能仅展平 C 就假定全部 backend 可用。未完成通用适配时，必须显式选择真实环境行宽的分轮模式，一轮每个真实行至多处理一个对应候选，其余安全占位；不能悄悄放松校验或伪造实例数。 + +### 5.2 碰撞世界与共享边界 + +同一 robot-relative 布局可共享不变的环境几何；arena offset 由适配器变换。不同 case、不同阶段动态物体、不同持物关系不能共用错误的 obstacle pose 或可变附件状态。 + +共享 world 只有在该规划批次所需碰撞状态一致时才成立；否则按正确来源映射 per-candidate world，或拆成兼容分组。复用 backend 时也要重新绑定当前世界，不能让另一个 case 的更新污染正在求解的批次。primitive 不直接构造 cuRobo/TOPPRA 专用 options。 + +关闭动态碰撞只可用于声明清楚的受限规划实验，不能代替正式数据必须通过的路径验证。把多个 grasp 放进“只返回一个 winner”的 goalset 也不等于获得多条候选轨迹。 + +### 5.3 按依赖复用，不重复完整流水线 + +| 改变内容 | 可尝试复用 | 必须重算或重验 | +|---|---|---| +| 只改运动时间 | 几何路径、兼容 IK、静态几何检查结果 | 时间网格、事件映射、动态限制与物理效果 | +| 改某个自由段路点 | 未受影响的目标/阶段与几何资源 | 该段 IK/路径、连接边界、依赖它的后缀 | +| 改 grasp/接触模式 | 场景几何与无关模板 | 接触段、持物关系及受影响下游规划 | +| 同 case 恢复初态或换兼容副本 | 满足同一初态契约的模板/几何候选 | 初态验证、坐标映射、runtime epoch、效果请求与碰撞绑定 | +| 外层切换 case | 真正场景无关的模板与机器人模型 | 场景相关路径、可行性与覆盖归属 | + +缓存键包含来源修订、场景内容、任务、机器人/标定、起点、阶段依赖和算子参数；计时计划再包含周期和执行限制。缓存设大小上限与淘汰策略，失效不清空其他 case 的有效数据。不能把旧 ActionPlan/session 当作可跨 reset 复用的值对象。 + +## 6. 覆盖驱动的候选池与执行调度 + +### 6.1 分场景集合选择与按需展开 + +对每个 case 分别统计阶段、接触家族、实际 IK 构型、EEF 位置/姿态、归一化关节状态、方向、速度和阶段转换状态。采用少量联合分桶及轨迹近邻距离，不构造完整高维网格。 + +质量门槛通过后，集合目标可写为 `Σ w_c · min(n_c, target_c)`；一个轨迹反复经过同一区域不重复刷分。几何按阶段进度对齐去重，速度/时长另算。同一路径的时间变体归入同一 geometry family。 + +调度使用估计的剩余价值，而非固定穷举： + +$$ +\mathrm{priority}(x) +=\frac{\widehat{P}(\mathrm{accept}\mid x)\, + \widehat{\Delta\mathrm{coverage}}(x)} + {\widehat{\mathrm{remaining\ cost}}(x)+\epsilon}. +$$ + +这是启发式优先级，不是最优性保证。保留探索预算和各目标 case/模式的最低配额，避免只采容易成功的区域。尚未校准的成功率使用保守先验，不把 grasp cost 直接当概率。 + +候选池不足时才继续展开接触、IK、几何及时间分支。可先对几何家族的一种时间方案回放，再根据价值与反馈决定是否增加其他时序；若失败可能由时序造成，也应保留有限替代尝试。不要为每条路径预先生成全部时间组合。 + +规划描述符只决定优先级。正式覆盖由通过验收且确认写入的实际轨迹更新；在途候选暂时预留覆盖额度，失败或写入失败释放。预算、独立探测和覆盖增益共同决定停止；只能报告在当前预算和采样器下趋于饱和，不能宣称穷尽连续可达空间。 + +### 6.2 场景与副本池解耦 + +`SceneCase` 是数据任务，`ExecutionSlot` 是物理资源。一个 case 可绑定多个副本；不同 case 的候选只在兼容槽中运行。副本共享准备信息和不变计算资源，不共享运行中的物理状态。 + +设 B 个物理槽，G 个活跃 case，每 case 分配 R_g 个副本，满足 `Σ R_g <= B`。例如 16 槽可用于 `1×16` 集中扩增一个场景，或 `4×4` 同时处理四个外部给定场景。副本数提升单场景并行度，但不意味着相同硬件的总吞吐量按倍数增长；最佳分配需实测规划、物理与渲染瓶颈。 + +两种池模式： + +- `per_env_case`：接收各行现有布局，各自维护 case；不要求它们相同，也不跨行直接使用轨迹。 +- `grouped_replicas`：宿主把外部给定 case 的初态准备到多个兼容槽；只在旧 episode 收尾后重绑定，拓扑/机器人不兼容时拒绝或分到独立宿主。 + +缺少同场景恢复/复制能力时，不得声称已批量验证同一场景的所有候选。可以只返回规划候选，或由宿主显式提供新 case 后重新规划；这些不是原 case 的剩余候选执行。 + +### 6.3 时长分桶、屏障与逐行补位 + +首版采用 `full_batch` 屏障，要求 Runner 独占被重置的整个宿主 batch，并优先将预计时长接近的候选安排在同一轮。完成行保持安全状态、停止增加训练帧；若效果需持续成立，提交前仍检查保持条件。真实任务等待与完成后的占位 hold 区分记录。 + +忽略其他开销时，全批有效步利用率约为 `Σ L_i / (B × max L_i)`。例如 200/220/240/800 步的一批仅约 46%；时长分桶可以减少此类浪费，但不改变任务必需时长。 + +逐行补位是后续能力：完成行验收后恢复并领取下一个兼容候选，其余行继续；所有动作仍由唯一宿主合并后 batched step。启用前需同时验证： + +- reset、速度清理、控制器、任务桥、观测历史和记录均行隔离。 +- 不因某行 settling 额外推进其他活动行；所有必要推进进入统一步进与记录。 +- 局部 reset 不清空其他行的任务状态，不使其碰撞绑定失效。 +- 每槽维持 candidate/attempt/epoch，终态证据在复用前冻结。 + +当前 [SimulationManager.update()](../../embodichain/lab/sim/sim_manager.py) 推进整个 world；部分 [随机化事件](../../embodichain/lab/gym/envs/managers/randomization/spatial.py) 会额外 sim.update，当前 [EmbodiedEnv.reset()](../../embodichain/lab/gym/envs/embodied_env.py) 还会清除活动任务桥。因此 `reset_ids` 存在并不证明全系统已支持安全的逐行补位。独立 world/process 分片可作为替代，但需计入显存与进程成本。 + +### 6.4 有界流水线 + +组织为 `提议/规划 → 待执行候选池 → 执行/验收 → 待写入队列`，各队列有条数与字节上限。达到高水位停止生产，低水位补充；writer 阻塞也须向上游反馈，不能无限持有图像和轨迹张量。 + +规划只读稳定快照，宿主独占物理步进；有状态 planner 的 world 更新与求解串行保护，或使用真正独立的 backend 实例。图捕获/warmup 在受控窗口完成，不与其他 CUDA 使用者无协调竞争。先重叠 CPU 整理/写盘与物理执行；同 GPU 的规划、物理和渲染是否并行由实测决定。 + +[Isaac Lab Mimic](https://isaac-sim.github.io/IsaacLab/v2.3.0/source/api/lab_mimic/isaaclab_mimic.datagen.html) 的按环境 action/reset 队列可作为职责分离参考，不意味着 EmbodiChain 可直接复用其调度器。 + +## 7. reset、初态恢复与单／多环境流程 + +### 7.1 明确区分三种状态操作 + +以下定义宿主协议的语义，不是普通 reset options；当前全批 `acquire_case / restore_initial` 已实现，checkpoint 仍为后续能力： + +| 操作 | 作用 | 候选与覆盖 | +|---|---|---| +| `acquire_case` | 宿主外层 reset 或读取外部指定场景，完成准备后注册 case | 按实际各行内容归属；新场景新建分区，不决定布局如何采样 | +| `restore_initial` | 为同 case 的下一条独立轨迹恢复指定初态，不重新采样布局 | 验证兼容后可消费剩余候选；重新绑定 runtime epoch | +| `resume_checkpoint`，后续 | 恢复任务中途的完整可续跑状态 | 只允许验证过的前缀/后缀衔接，不等同于 episode 初态重建 | + +固定 seed 不代表不同子环境布局相同，也不能替代初态一致性检查。复用初态需显式准备策略，保留必要的机器人/控制器初始化；不能偷偷跳过全部 events。startup/reset/interval 中会改变本 case 固定条件的随机项必须与宿主约定处理。 + +首版 `restore_initial` 可采用“完整 episode 重置 + 按 case 重建指定初态 + 校验”，无需先实现任意接触状态 checkpoint。恢复内容至少包含任务所需的机器人/物体状态与速度、夹爪/约束、控制器和 task/manager 初始状态，并完成 settling。已有 [轨迹状态读写](../../embodichain/lab/gym/utils/trajectory_state.py) 只是其中一部分。 + +当前已落地 [FixedSceneHost / InitialStateProfile](../../embodichain/lab/trajectory_generation/initial_state.py) 与 [SimInitialStateAdapter](../../embodichain/lab/trajectory_generation/integrations/sim.py) 的全批初态基础：一个配置所有的固定基座 robot 和全部普通刚体，保存独立物理状态并在恢复前预检，使用可信 profile 准备任务/控制器、标识固定条件和验证初态。Gym 的 generation lease/prepare 入口不调用普通 reset，且与纯 sim host 共用 simulator batch 所有权；调用者先冻结旧 episode，再丢弃旧录制、准备、恢复、settle、standard episode manager reset、verify，验证后刷新首观测并播种记录。每次准备或释放使旧 epoch binding 失效，失败不允许继续执行。已通过物理关闭/开启两种模式的真实 headless CPU backend 状态恢复 smoke，并接入手写 qpos Runner 与真实 UR5 free-motion profile。固定条件与控制历史不由物理快照自动覆盖；PickUp profile、接触与长时间 settling 验收仍需后续完成。 + +这种“随机 reset 与指定状态 reset 分开”的接口可参考 [Isaac Lab reset_to](https://isaac-sim.github.io/IsaacLab/develop/source/api/lab/isaaclab.envs.html#isaaclab.envs.ManagerBasedEnv.reset_to)，不能据此推断现有项目已支持。 + +### 7.2 统一采集流程 + +```mermaid +flowchart TD + A["宿主提供一个或多个 case
完成初始化并取得快照"] --> B["各 case 按需规划、筛选
维护独立候选池"] + B --> C{"执行槽分配"} + C -->|"单环境"| D["每次执行一条兼容候选"] + C -->|"多环境"| E["按 case 分组
各副本执行不同候选"] + D --> F["验收、冻结证据
提交合格数据"] + E --> F + F -->|"原 case 继续"| G["恢复各槽所属 case 初态
刷新观测与运行时绑定"] + G --> B + F -->|"宿主显式切换 case"| A + F -->|"目标达到或预算耗尽"| H["排空在途写入
汇总实际完成与覆盖"] +``` + +回到候选池时优先取兼容缓存，不重新生成整个池。多环境首版按第 6.3 节的全批屏障收尾，图中的逐行验收不表示已经允许逐行 reset。 + +| 事件 | 下一步 | +|---|---| +| 纯采样/IK/规划失败且宿主未变化 | 不 reset，只更新 audit、消耗规划预算并继续采样 | +| 候选完成、失败、超时或受控基准采集结束 | 冻结验收与数据，再恢复初态；失败不进入专家数据 | +| 中间阶段完成 | 不 reset，从实际边界状态规划/执行下一阶段 | +| 初态恢复后不匹配 | 拒绝执行并重试准备；仍不匹配报告错误，不静默当作新 case | +| 宿主显式提供新布局/任务/标定 | 重新注册或匹配 case，失效不兼容计划；旧覆盖保留在原分区 | +| 预算耗尽/取消 | 停止接收新工作，安全收尾在途执行和写入，不降低验收标准 | + +Job 的预算、ID 序列和各 case 历史跨 reset 保留；episode 的控制游标、观测历史和记录在结束后重置。队列能否复用取决于 case/初态兼容性，不是“跨 reset 保留”就永远可执行。 + +### 7.3 起点、观测与准备顺序 + +全部初始化、复制和 settling 必须完成在新 episode 首帧播种之前。若准备发生在 env.reset 返回后，适配器需显式暂停准备段记录，刷新观测、重新播种首帧，再允许执行；不能 set_qpos 后继续使用旧观测。 + +启用 start_state 时，只选择部署允许的起点。若部署始终从固定 qpos 开始，需真实执行合法连接前缀；不能把人为设置的中途状态当作策略自然起点。准备动作不作为专家标签，除非它本身是要学习且完整记录的任务前缀。 + +采集期间抑制隐式 auto-reset，沿用 [demo 执行器](../../embodichain/lab/gym/envs/demo.py) 由外部 Collector 掌管收尾的约定。任务终态验收必须发生在状态被重置前。 + +## 8. 验证、真实数据与提交 + +### 8.1 统一验证链 + +| 层 | 必检内容 | +|---|---| +| 输入与语义 | 有限 SE(3)、工具标定、合法因子、阶段顺序和锚点 | +| 运动学 | IK/FK residual、关节限位、实际分支及连续性 | +| 路径 | 自碰撞、环境/双臂碰撞、携带物体扫掠、阶段接触许可及段间运动 | +| 时间与控制 | 周期、速度/加速度及任务要求的其他动态限制、事件顺序 | +| 物理与任务 | 跟踪误差、接触建立、滑移/掉落、阶段效果和完整任务成功 | +| 数据质量 | 合理路径/时长、去重、观测动作因果对齐和谱系完整性 | + +使用连续碰撞检测或与运动尺度适配的路径采样，不能只检查关键帧。允许手指在接触阶段接触目标，不意味着整段忽略目标；持物阶段检查对象随候选运动的几何。 + +每项记录 `not_run / passed / failed / unavailable`、失败阶段及连续度量。缺少必需检查的候选只能留在规划库；规划成功、效果投影和命令完成均不等于物理成功。原始“专家”参考也需验收；质量 ratio 与同起点、兼容任务/模式的基准比较。 + +### 8.2 记录实际因果序列 + +记录 `observation_t → action_t → observation_{t+1}`，包含时间、实际机器人/接触状态及阶段。明确源动作、训练标签和下发命令的表示与转换 profile；不能将计划 qpos 或状态差自动当作动作。 + +每个执行分支重新采集图像；不混用其他分支或另一次执行的观测。真实等待计入实际时长，padding 和完成后的占位不成为训练帧；计划时长与实际时长分别记录。[LeRobotRecorder](../../embodichain/lab/gym/envs/managers/datasets.py) 已有因果配对通路可复用。 + +### 8.3 Gym reset 的保存顺序 + +当前同步持久化基础位于 [LeRobotEpisodeSink](../../embodichain/lab/trajectory_generation/sinks.py)：每个已验收 `ExpertEpisode` 写入独立 LeRobot shard，T 个训练帧与末观测/T+1 实测时间 sidecar 一同封存并读回，谱系/验证 metadata 和 manifest 也必须读回成功才给 confirmed receipt。相同 commit ID 的相同 payload 可在当前进程内幂等重试；已确认提交只读核验、不重写。输入或时钟/容量不支持时拒绝，写入失败返回失败回执。该 sink 无活环境引用；`QposRolloutExecutor` 已负责冻结实际命令与观测，`GenerationRunner` 接入实测验收、Session 和保存回执。manifest 未列出的片不能作为正式专家数据消费。 + +当前 [BaseEnv.reset](../../embodichain/lab/gym/envs/base_env.py) 与 [EmbodiedEnv](../../embodichain/lab/gym/envs/embodied_env.py) 的顺序为： + +```text +Collector 先完成旧 episode 验收并冻结终态证据 +→ BaseEnv 缓存旧任务成功状态 +→ sim.reset_objects_state(reset_ids) +→ DatasetManager 消费旧 buffer +→ 清空旧记录和 metadata +→ reset events 与 managers +→ 获取新观测，EmbodiedEnv 播种首帧 +``` + +物理状态先于 DatasetManager 保存被重置，因此保存阶段不得重新读取现场来判定旧 episode。异步写入需在 buffer 清空前取得独立 payload。 + +现有全量 reset 的选择性 dataset 提交可写为： + +```python +# accepted_ids 来自 reset 前完成的任务、质量与去重验收。 +env.reset(options={"save_data": False, "commit_env_ids": accepted_ids}) +``` + +这仅展示已有提交语义，**不等于实现 restore_initial**。该调用仍可能运行布局随机化；新适配器必须把同 case 准备与首帧播种按第 7 节接好，不能直接执行旧候选。 + +commit_env_ids 是被 reset 行的无重复子集，只指定 dataset 保存行，不指定物理 reset 范围；显式提交会绕过默认成功过滤，调用者必须先验收。save_data=False 下其他 camera/trajectory 自动产物需要单独明确提交策略。[现有收集入口](../../embodichain/lab/scripts/run_env.py) + +### 8.4 提交确认与训练使用 + +分别维护 `proposed / planned_valid / rollout_attempted / validated_accepted / pending_write / committed`。只有实际动作才计 rollout；只有收到持久化确认才增加 committed 和正式覆盖。回执按 episode/commit ID 去重并归属原候选，重试写入不重复计数，不能按当前 slot 绑定解释旧回执。最后一批先预留剩余配额，按稳定候选顺序提交，避免超额；失败不靠放宽门槛凑数。 + +Gym 可复用 DatasetManager 及 [异步 recorder](../../embodichain/lab/gym/envs/managers/async_datasets.py)，但需向会话补齐结果确认。直接 sim 使用宿主无关 EpisodeSink，或抽取通用写入器，不伪造现有 Gym recorder 所需 env。Runner 负责 finalize/drain 与错误回传，不依赖 env.close 隐式提交当前 episode。 + +训练按轨迹家族和父分支分组，时间变体不跨 train/validation；限制近重复、公共前缀和慢速版本的采样权重。完整任务与成功片段分别标识。多模态动作宜使用可表达模式且具时间一致性的策略，或部署时可得的条件；candidate ID、规划 seed 和未来成功标志不能作策略输入。[Diffusion Policy](https://diffusion-policy.cs.columbia.edu/) + +## 9. 超参配置与可复现性 + +### 9.1 配置所有权 + +公共 `TrajectoryAugmentationCfg` 管理因子、去重和覆盖；`TrajectoryGenerationJobCfg` 包含前者，并管理来源、规划预算、池调度、reset、验证和 sink。均使用项目 `@configclass`。 + +采用独立 `generation.yaml`，由启动器显式加载；不向现有 env.yaml、Task Program 或 EventManager 暗加语义。模板与执行 profile 拥有硬约束，job 只能收窄范围。长度用米、角度用弧度、时间用秒；关节偏移按各关节有效范围归一化。control_dt 从宿主取得，不由 job 静默覆盖。 + +### 9.2 配置示例 + +以下为**待实现 schema**，不能直接交给当前 config_to_cfg。示例输入是阶段标注齐全的手写 EEF 模板；数值是启动实验值，不是验证过的最优默认值。 + +```yaml +source: + kind: handwritten + source_id: fixed_pick_place_waypoints + template_id: reference_0 + +augmentation: + seed: 20260905 + start_state: {mode: provided} + factors: + contact: {enabled: false} + ik: + enabled: true + max_seeds: 8 + max_solutions_per_target: 4 + allow_branch_switch: false + approach: {enabled: false} + spatial: + enabled: true + method: via_points + variants_per_branch: 2 + max_via_points: 2 + position_offset_norm_m: [0.0, 0.03] + rotation_offset_norm_rad: [0.0, 0.15] + preserve_phase_endpoints: true + timing: {enabled: true, duration_scales: [0.8, 1.0, 1.2]} + contact_timing: {enabled: false} + recovery: {enabled: false} + phase_overrides: + transit: {operators: [via_points, retime]} + approach: {operators: [retime]} + grasp: {operators: []} + transfer: {operators: [via_points, retime]} + release: {operators: []} + coverage: + descriptor: phase_joint_eef + selection: value_per_cost + position_bin_m: 0.03 + rotation_bin_rad: 0.17 + joint_dedup_normalized_tol: 0.01 + geometry_samples_per_phase: 32 + target_per_cell: 4 + exploration_fraction: 0.2 + saturation_min_proposals: 512 + saturation_patience_rounds: 8 + min_new_cells_per_round: 1 + +planning: + batch_mode: candidate_buckets + batch_buckets: [16, 32] + cache_budget_mb: 1024 + +execution: + pool_mode: grouped_replicas + rollout_slots: 4 + active_case_limit: 1 + replicas_per_case: 4 + scheduler: full_batch + duration_bucket_edges_s: [2.0, 4.0, 8.0] + ready_low_watermark: 8 + ready_high_watermark: 32 + ready_max_bytes: 268435456 + overlap_planning_and_physics: false + +reset: + outer_mode: provided + inner_mode: restore_initial + prepare_profile_id: fixed_scene_initial_state + initial_state_tolerances_profile_id: fixed_scene_tolerances + on_initial_state_mismatch: error + +validation: + validator_id: fixed_pick_place_success + profile_id: fixed_pick_place_verified + motion_limits_profile_id: robot_execution_limits + require_path_collision: true + require_task_success: true + path_length_ratio_max: 1.5 + duration_ratio_max: 1.5 + +collection: + target_committed_episodes: 256 + max_proposals: 8192 + max_rollout_attempts: 2048 + max_attempts_per_candidate: 1 + max_wall_time_s: 3600 + +persistence: + sink: lerobot + accepted_only: true + async_write: false + pending_episode_limit: 8 + pending_max_bytes: 536870912 + save_audit: true + split_unit: trajectory_family +``` + +配置语义： + +- phase_overrides 只控制许可的局部几何/时间算子；不能启用全局关闭的因子或放宽模板。grasp 的空 operators 不妨碍对固定接触目标求 IK。 +- 偏移在模板声明坐标系内采样；duration_scale > 1 表示延长允许 retime 的阶段，不自动缩放接触等待。 +- 原子来源可开启有声明的 contact；纯 qpos 来源默认不开 IK，除非显式支持 FK 后重新求解。两类来源共享后续配置。 +- `candidate_buckets` 必须有第 5.1 节的候选适配；受限 backend 需显式选 `env_rows`，不能忽略配置。grouped_replicas 必须通过初态复制/恢复能力检查。 +- `outer_mode: provided` 接收调用者给定 case，不自行触发布局采样。若选择 `host_reset`，环境 RNG/布局仍由宿主拥有，每次结果重新归属 case；内层 restore 不推进该布局随机流。 +- 时长桶最后含超出最大阈值的溢出桶；水位、缓存与字节限制同时约束内存。单个请求/episode 超限时明确拒绝或使用预先声明的流式 sink，不突破上限。 +- 多 case 时增加 active_case_limit 并约束总副本数不超过可用槽；各 case 的覆盖、去重独立，job 目标可全局累计，但不能用其他 case 完成声明的局部配额。 + +### 9.3 校验、随机流与调度重现 + +加载时拒绝未知字段、非法范围/单位、未注册 source/validator/profile、无效 phase、非正预算、不递增桶边界、不合法水位及超过宿主槽数的副本配置。初态 profile 明确实体位置/旋转、机器人关节/速度和控制器检查；运动 profile 明确每关节限制。仅有名称没有定义时不能启动。 + +使用显式局部 Generator 和固定 hash 算法派生种子： + +```text +H(job_seed, source_id, source_revision, scene_case_id, + candidate_ordinal, operator_id, attempt_id) +``` + +不使用 Python 进程随机 hash、物理 env ID 或全局环境 RNG 作为候选身份。纯因子提议应不受 slot 分配与 chunk 改变影响；backend 批大小可能影响数值求解，因此记录实际桶、版本、参数及输出，不承诺逐位一致。 + +默认按逻辑轮次和 candidate ID 合并反馈，便于对照实验。将来吞吐优先异步模式可按完成顺序更新，但必须记录调度事件用于审计，不能同时承诺与同步模式完全相同的候选序列。预算/历史不因 restore 清零。 + +## 10. 可选的高级优化 + +| 模式 | 适用条件 | 必须守住的边界 | +|---|---|---| +| 公共规划前缀缓存 | 多候选的前段约束与起始状态相同 | 复用计算不等于复用物理成功证据 | +| 检查点分叉执行 | 前缀长、分支发生晚，且完整状态可恢复 | 物理、控制器、任务、接触/约束及观测历史一致；验证真实前缀/后缀连接，不只恢复 qpos | +| 低成本物理预筛，再完整采集 | 失败率高且渲染占主要成本 | 计入二次执行成本；正式采集重新验收，图像与动作必须来自同一次正式执行 | +| 实际偏离状态的专家恢复 | 需要闭环纠偏且有可靠后缀专家 | 通过真实前缀到达偏离状态；致偏离动作不标为专家动作，恢复边界不混入连续训练窗口 | + +[DemoExecutionCfg](../../embodichain/lab/gym/envs/demo.py) 当前明确未提供 checkpoint/resume 协议；片段保存或轨迹状态回放不代表该能力。首版不依赖上述高级模式，尤其不把刚性初态重建和任意接触中途恢复混为一谈。 + +若公共前缀被多个后缀引用，数据存储和采样要记录共享谱系，避免把重复前缀算作新增独立覆盖。恢复数据可针对学习策略真实访问的状态追加专家监督，借鉴 [DAgger](https://proceedings.mlr.press/v15/ross11a.html) 的分布适配思路；任意随机扰动不具有相同保证。 + +## 11. 实现缺口、交付顺序与验证 + +### 11.1 现有基础与待补模块 + +以下记录当前实现与剩余设计范围。真实验证已包括初态恢复、Panda 动态障碍采样检查，以及普通重力下 UR5 自由运动的执行、实测验收、LeRobot 写入/读回。Panda 手指漂移负例被锁定关节模型检查拒绝。真实 Runner 验证为 B=1、direct-sim、20 个命令，不代表 PickUp、真实 Gym 采集或四组合 M1 已完成；详细命令和验证记录见[实施计划](fixed_scene_expert_trajectory_augmentation_implementation_plan.md)。 + +| 模块 | 已有基础 | 仍需补齐 | +|---|---|---| +| 公共核心与来源 | qpos 模板/候选/配置、严格能力校验、demo 候选输入、原子初始计划供应入口 | EEF via-point 因子、完整原子来源导出与候选消费 | +| IK 与规划 | MotionGenerator、cuRobo、真实 env_rows 分轮、EEF 显式样本 IK 与已解分支/FK 保留 | 多解枚举/去重、主动绕障路线生成与有界修复、独立候选容量桶与更广 backend 能力 | +| 路径验证 | free/no-held qpos 加密采样、实测路径重验、运动限值与质量独立 gate | 接触/持物/夹爪变化语义、连续碰撞与任务级 PickUp 验收 | +| SceneCase 与副本池 | 全批物理初态捕获/恢复/校验、可信 profile、独占 host/epoch、Gym 观测重播种 | 同 case 副本重绑定、接触任务 profile 与长期 settling 验收 | +| 会话与调度 | 按 case 分区、局部 RNG、几何去重/覆盖、提议/执行/写入配额及条数/字节上限 | 依赖缓存、时长调度、有界异步 pipeline 与成本调度 | +| sim/Gym 数据闭环 | 手写 qpos Runner、实际命令/观测冻结、同步 LeRobot 封存/读回/确认，真实 UR5 direct-sim 正例 | 真实 Gym/PickUp、原子来源、统一注册/CLI 及四组合 M1 验收 | +| 逐行异步／分叉 | 部分行级接口、自然片段保存 | 全链路行隔离、独立任务上下文、完整 checkpoint 及连续性协议 | + +### 11.2 交付顺序 + +1. **正确闭环**：公共契约、case/slot 身份、合法初态准备与校验、手写 EEF 模板和 PickUp、sim/Gym 执行记录；全批屏障、明确失败与提交确认。 +2. **优先效率项**：同 case 副本池、候选适配、稳定容量桶、依赖缓存、时长分桶、有界流水线与覆盖预算；与第一步做同预算对照。 +3. **扩大运动覆盖**：多 grasp/多 IK 分支、手写多段和 PickUp → Place，再推广 AxisAlign、Slide、OpenDoor、HandOver 的特定约束。 +4. **按瓶颈选高级项**：行隔离成熟后逐行补位；长公共前缀再考虑分叉；需要闭环泛化时加入真实恢复数据。 + +首版必须验证四个组合：原子 × sim、原子 × Gym、手写 × sim、手写 × Gym。复用核心不要求四条路径使用相同执行封装。 + +### 11.3 聚焦回归 + +| 测试面 | 关键用例 | +|---|---| +| 核心边界 | 算法不直接导入 Gym/backend，motion 父包按需加载；缺标注、非法控制表示、无能力配置明确失败 | +| 候选正确性 | 空 grasp、混合 IK 失败、失败 seed 隔离、实际分支保持/去重、索引 round-trip | +| 避障扩增 | 固定场景与锚点、不同有效绕行路线、修复后路线去重、最终输出重验、预算耗尽及能力缺失拒绝 | +| 批量规划 | C ≠ B、固定桶 padding、受限 backend、重复来源行、根坐标与碰撞世界不串用 | +| 场景与恢复 | 多行不同布局、同 case 多副本、arena offset、恢复失配拒绝、旧 epoch 不执行 | +| 时间与接触 | 变长 timing/事件一致、时间变体不刷空间覆盖、携物碰撞与阶段接触许可 | +| 宿主与记录 | 唯一步进、ControllerAction 不重复处理、准备后首帧刷新、占位不写训练帧 | +| 提交与队列 | reset 前证据冻结、仅提交合格行、writer 失败不计 committed、队列满反压、尾批不超额 | +| 并发隔离 | planner world 绑定互不污染；一行 reset/清速度/清桥不影响其他行；失败时禁用异步模式 | +| 学习兼容 | 规划成功物理失败不入专家集，家族 split 无近重复泄漏，恢复/片段边界清楚 | +| 旧行为 | 默认在线逐环境单 winner、唯一 env IDs、ActionPlan 效果与恢复协议不变 | + +聚焦验证可从 [Atomic Actions](../../tests/sim/atomic_actions/)、[MotionGenerator](../../tests/sim/motion/planners/test_motion_generator_batched.py)、[BasePlanner](../../tests/sim/motion/planners/test_base_planner.py)、[Demo 执行](../../tests/gym/envs/test_demo.py)、[Dataset recorder](../../tests/gym/envs/managers/test_dataset_functors.py) 扩展。纯张量单测之后，再做无 GUI 的真实执行验证。 + +### 11.4 效率与学习效果验收 + +在同一组外部给定 case、相同副本总数、相同传感器/控制周期、质量门槛和每 case 数据目标下，逐项比较：基础全批 → 缓存/容量桶 → 同场景分组/时长分桶 → 流水线 → 可选逐行补位。不要用改变 case 数量或降低验收标准制造吞吐提升。 + +记录规划有效率、物理接受率、几何/时间去重率、committed/min、覆盖增益/min、首次达到覆盖目标的时间；同时拆分 reset/settling、IK/规划、物理、渲染、拷贝/写盘、屏障等待、warmup 和显存峰值。固定形状与最小压紧、单 GPU 串行与重叠都需做消融，先优化实测瓶颈。 + +最终在相同数据预算下比较模仿学习成功率，并保留未见过的起点/运动参数组合进行评估。固定场景实验只能证明相应运动分布上的收益，不据此宣称布局泛化或连续可达空间已完全覆盖。 diff --git a/docs/design/fixed_scene_expert_trajectory_augmentation_implementation_plan.md b/docs/design/fixed_scene_expert_trajectory_augmentation_implementation_plan.md new file mode 100644 index 000000000..11065735c --- /dev/null +++ b/docs/design/fixed_scene_expert_trajectory_augmentation_implementation_plan.md @@ -0,0 +1,351 @@ +# 固定场景专家轨迹扩增：实施计划 + +- 状态：实施中；手写自由运动与离线原子 PickUp 的 direct-sim 采集闭环已落地，完整 Atomic Runtime/Gym/M1 仍待实现与验收。 +- 依据：[设计文档](fixed_scene_expert_trajectory_augmentation_design.md)，2026-09-05。 +- 代码核对基线：`fb7228e2`。首批实现基于该版本，已运行 CPU 行为测试；尚未完成四组合真实仿真和性能验收。 +- 交付原则：先完成可验收、可持久化的四条执行路径，再优化吞吐，最后扩大运动覆盖。 + +## 当前实施进度（2026-09-05） + +| 计划项 | 已实现 | 尚待实现与验收 | +|---|---|---| +| PR 1 | `lab.sim.motion.trajectory_augmentation`：qpos 值对象、nested 配置严格校验、GenerationSession、局部 RNG、候选与提交身份、条数/字节预算、实际关节几何去重与写入配额 | EEF 原始模板和宿主能力注册装配随实际适配器接入 | +| PR 2 的初态基础 | 全批初态复制、恢复、验证与独占 host；Gym 准备/重播种；UR5 自由运动 profile；新增四行 PickUp profile 和抓取结束后的整批循环恢复 | 任意接触 checkpoint、逐行恢复和复杂 settling 仍不支持；Gym PickUp 待验收 | +| PR 4 的基础算子与自由段适配 | 显式自由段 `joint_residual`、retime 与阶段索引重映射、采样速度/加速度检查；`EEFPath / EnvRowMotionPlanner` 的真实 env_rows 分轮、显式 EEF 样本 IK、已解分支保留/FK 验证、自由 qpos 分段加密碰撞检查 | EEF via-point 因子生成、接触/持物/夹爪变化的完整碰撞语义与真实 backend 任务验收；离散路径检查不代表连续碰撞证明或物理成功 | +| PR 5 的手写 qpos/free-motion 闭环 | `GenerationRunner / MotionLimitsProfile / QposRolloutExecutor` 连接 profile 身份、全批初态、sim/Gym 实际命令/观测冻结、规划与实测验收、Session 配额和 sink 确认；真实 UR5 direct-sim 正例 committed 1，Panda 锁定关节漂移负例 committed 0 | 手写 EEF/PickUp、真实 Gym 任务采集及多行/多 episode 任务验收；当前结果不代表 PR 5 全部原定验收或 M1 完成 | +| PR 6 的候选输入与离线模板 | 既有 `initial_plan_provider` 入口；新增 `export_pickup_templates`，从 MoveEndEffector → PickUp 离线编译导出受保护阶段、mimic 几何和实测保持段；接入 Runner 的离线 qpos 采集 | 候选到新的运行时 ActionPlan 的适配、Atomic Runtime tracking/recovery 和 Task Program/Gym 四组合接线；离线导出不代表这些完成 | +| 新增 PickUp 接触闭环 | `PickUpMotionValidator`：完整 URDF 碰撞凸包、mimic 与持物几何、分阶段接触规则、CPU 物理子步接触证据、抬升到保持的相对漂移和双指接触验收；`cube_pickup_collection.py` 四行多轮采集，保存视频和确认后的 LeRobot 数据 | 仅固定基座未缩放 URDF + 盒状刚体 + CPU 物理；通用 mesh 场景、连续碰撞、Gym 子步接口另行扩展 | +| PR 3 的同步保存基础 | `ExpertEpisode` 与 `CommitReceipt`；Session 预留/确认/失败重试；LeRobot 分片封存/真实回读；新增数值矩阵展开及原始 shape/终态保留；确认提交只读幂等核验 | 重启恢复、异步 pipeline 等后续范围；完整四组合真实采集仍需 PR 6 验收 | + +当前 `restore_initial` 使用显式全批物理状态适配与可信 profile，不调用普通 reset；同步 sink 只有封存并读回真实文件才确认提交。手写 qpos 的宿主证据冻结、实际控制标签、Runner/Session/sink 已接通，并通过普通重力下 UR5 自由运动真实采集。新增 direct-sim PickUp 的受限接触/持物检查和任务 profile；统一配置注册/CLI、完整原子运行时来源和 M1 四组合发布门槛仍未完成。当前示例显式构造可信服务；配置 ID 本身不代表服务已经注册。 + +新增 API、测试和现有行为变更已同步对应文档及 agent context。下一步接入 Atomic Runtime 的候选消费与等价物理子步证据，再推进 Gym PickUp、via-point 因子和统一配置启动器；当前离线回放不扩大为完整原子运行时或 M1 四组合能力声明。 + +模块归属 review 后，`solvers`、`planners`、`workspace` 与 `trajectory_augmentation` 统一归入 `embodichain.lab.sim.motion`。公共 API、测试、示例和文档使用新路径，不提供旧路径兼容包。`motion` 父包仅按需加载子包；扩增算法保持无直接 Gym/仿真 backend 依赖，但其公共导入不再承诺绕过 `lab/sim` 初始化。该调整不改变下述 M0/M1 验收要求或未完成状态。 + +首批验证记录：核心、原子动作与 Gym/Task Program 组合回归 `1006 passed, 1 skipped, 3 deselected`；补完前置写入容量预留和数值边界后，最终核心测试 `138 passed`。全仓 Black 检查通过，API 文档覆盖 `1718/1718`。Sphinx dummy 构建成功，本次新增模块无相关告警；其他模块仍有文档告警。这些是代码契约验证，不替代 M0/M1 的真实采集验收。 + +motion 重组回归：`1246 passed, 113 deselected`（排除 requires_sim/gpu/slow），独立进程的新导入和配置解析测试 `3 passed`。PR 2 物理适配 CPU 契约测试 `35 passed`；真实 headless CPU backend 的物理关闭/开启两种模式 smoke `2 passed`，以单关节 robot 与 cube 检查当前状态/drive target/速度/effort 的捕获和恢复，不覆盖接触抓取或长时间 settling。Gym 与宿主集成测试 `119 passed`；相邻 `root_articulation` getter 的 `root_link_name` 修复及回归中，no_sim 集合 `7 passed`。这些集合分别记录，不相加为去重总数；参考 UR5 PickUp 的四组合真实采集验收仍未完成。 + +PR 3 同步 sink 测试 `17 passed`：真实 LeRobot 数值/RGB/scalar 与 float64 精度读回、终态/T+1 时间 sidecar、writer/evidence/manifest/verify 失败重试、partial 重建、confirmed duplicate 只读以及时钟/容量拒绝。PR 4 首轮 env_rows/IK/碰撞适配 CPU 契约测试 `14 passed`；扩展后的 planning + cuRobo planner CPU 回归为 `82 passed, 1 skipped, 3 deselected`，覆盖分轮 scratch、float64 保留、锁定关节 cache 失效及实测路径重验。真实单 Panda、CPU physics + CUDA cuRobo 0.8.0 smoke `1 passed`：相同 backend 对远处动态方块的短 free 路径返回 passed，将方块移至 TCP 后返回 failed。这不包含接触抓取、持物扫掠或 M1 四组合验收。 + +PR 5 组合回归 `1379 passed, 1 skipped, 118 deselected`；补充有界 `last_failures` 审计后，Runner/executor 联合测试 `70 passed, 1 skipped, 1 deselected`。真实物理开启的 CPU 单关节 robot + static cube 执行器 smoke `1 passed`；真实 Runner 测试 `2 passed`，分别验证 UR5 接受并确认 1 条、Panda 漂移拒绝且提交 0 条。示例命令独立运行 exit 0，UR5 在 1 秒内记录 21 个实测观测/20 个命令，全部 gate 通过并读回 LeRobot。全仓 Black 检查通过（867 个文件）。以上集合有重叠，不相加；真实验证范围仍为单行 direct-sim free motion。 + +上一轮文档校验：公共导出覆盖 `1751/1751`，checker 测试 `8 passed`，Sphinx dummy 构建成功。本次新增 API/指南没有相关告警；构建仍报告其他既有文档的 46 条告警。`git diff --check -- docs` 通过。 + +### 本轮 PickUp 实施范围 + +- 来源通过 `AtomicActionEngine.compile` 复用现有 MoveEndEffector/PickUp 规划；`export_pickup_templates` 保留接近/闭合/抬升边界，展开 mimic 坐标，并加入 30 个真实保持命令。只增强 transit，不改变接触时序。 +- 碰撞适配使用 URDF **碰撞**形状的保守凸包和完整关节 FK，包含指尖、mimic、随 TCP 移动的目标盒体和隐式地面；实测重验使用真实目标位姿。两条关节边以内的结构自碰撞排除沿用 cuRobo 规则。验证有界采样，不声明连续碰撞保证。 +- CPU 每个 5 ms 子步读取原生接触，限定指尖/目标/支撑/固定安装部位的合法组合与接近段进入距离，拒绝未知/跨行/非法/过深接触和数据容量溢出。保持阶段双指真实接触比例至少 95%，抬升至少 12 cm；从抬升到保持检查 TCP 相对位置/转角漂移，避免把掉落或滑移数据提交。 +- 例子固定机械臂刚度 200000、手指打开目标距下限 1 mm，保留 0.05/0.08 rad 的末端/路径跟踪门槛。普通重力下并行执行 0°/90° cube 抓取，整批恢复后继续；达到 confirmed committed 目标才成功退出。 +- LeRobot 新增数值矩阵支持，训练帧按 C order 展开为向量，`observation_shapes` 保存原始形状，终态仍是矩阵；`commanded_joint_indices` 标记实际下发的 active 列。每轮视频包含真实初态与终态，拒绝的尝试也保留在视频和审计中。 +- 本轮没有消费 `initial_plan_provider`、提交预测符号效果或替代 Atomic Runtime 的 tracking/recovery。后续需将这些运行时契约与新增物理验收接通，才能完成 PR 6 的完整目标及 M1 四组合。 + +本轮同步 main 的显式物理步进 API 后，motion、原子动作、采集、Gym 初态与 demo、仿真管理器及 gizmo 组合回归为 `1542 passed, 1 skipped, 168 deselected`；真实并行采集、异常退出清理和原有视频示例为 `3 passed, 1 deselected`。四行两轮共 proposed/attempted/committed `8/8/8`，每条保持 1.5 秒、双指接触覆盖 100%，最小抬升 17.90–17.92 cm，相对位置最大漂移 0.35–0.73 mm；LeRobot 和 544 帧 H.264 视频均已真实回读。全仓 Black、`git diff --check` 通过；API 覆盖 `1761/1761`、checker 测试 `8 passed`，Sphinx dummy 构建完成，仍有文档告警。这些测试集与前述记录有重叠，不相加。 + +当前已实现部分合并为一个 draft PR 提交 review。下文的 PR 编号保留为实施计划的工作划分，不表示这些 PR 已独立创建或全部验收。 + +## 1. 实施范围与里程碑 + +| 项目 | 首阶段 M1 | 后续阶段 | +|---|---|---| +| 参考任务 | 单臂、刚体 cube 的 PickUp：接近、闭合、抬升及保持；手写 EEF 和原子来源使用同一任务条件 | PickUp → Place、手写多段及其他技能 | +| 参考 embodiment | 建议复用 UR5 + DH PGI 配置，先实测基准与碰撞能力 | Franka 用于第二 embodiment 回归；双臂独立扩展 | +| 来源 × 宿主 | 手写 × sim、手写 × Gym、原子 × sim、原子 × Gym，四个组合全部通过 | 参数化工厂和更复杂的 Task Program | +| 扩增 | 已标注自由段的少量 via points、合法 retime；固定接触目标，先保留一条有效 IK 分支 | 多 grasp、多 IK、approach、contact timing | +| 起点与 reset | `provided` 起点；无中途接触约束的 episode 初态重建、校验；布局由宿主提供 | 部署允许的多起点；checkpoint 与 recovery 另立项 | +| 执行与规划 | 先 `B=1`，再多行不同 case；`full_batch` + 显式 `env_rows` 分轮 | 同 case 副本、`candidate_buckets`、时长调度 | +| 保存 | 同步 sink、逐提交确认、仅接受合格 episode；从第一版限制队列条数和字节 | 有界异步写入及 CPU 整理与执行重叠 | + +M1 的 PickUp 成功表示完成所声明的抓取任务，不表示已经实现完整 pick-and-place。现有 repeated_pick_place 部署可提供场景与 embodiment 参考，但需要单独准备验收 fixture，不能把现有 open-loop 执行结果直接视为专家质量基准。 + +设计第 9.2 节的完整 YAML 包含 M2/M3 能力。实施时分别提供各里程碑可运行的 `generation.yaml`；未实现的模式和启用因子必须在启动前报错，不能静默降级。 + +## 2. 实施前代码基线与实施影响 + +| 已核对事实 | 实施影响与主要落点 | +|---|---| +| `lab/__init__.py` 会导入 sim/Gym 等子包 | 公共核心归属 `lab/sim/motion/trajectory_augmentation/`；验证算法无直接 Gym/backend 依赖，以及 motion 父包按需加载的边界 | +| `configclass.validate()` 主要检查 `MISSING` | 配置加载还需显式校验未知字段、范围、单位、交叉约束和宿主能力；不修改全局 configclass 语义 | +| [PickUp](../../embodichain/lab/sim/atomic_actions/primitives/pick_up.py) 提前筛掉 grasp 分支，预筛 IK 结果未直接用于后续轨迹 | 候选保留与已解关节目标传递需要新增接口；空 grasp 行也要覆盖 | +| [BasePlanner](../../embodichain/lab/sim/motion/planners/base_planner.py) 批量校验绑定 `robot.num_instances` | 首版使用 `env_rows`；C 与 B 解耦由独立适配层完成，不能仅展平张量 | +| [ExecutionSession](../../embodichain/lab/sim/atomic_actions/execution.py) 启动时重新规划 | 需要正式的候选消费入口，并保留既有 tracking、效果验证和 recovery 协议 | +| [execute_demo_episode](../../embodichain/lab/gym/envs/demo.py) 没有显式候选 segments 参数 | 新增明确输入，保留旧 hook 调用；不临时替换环境方法 | +| [BaseEnv.reset](../../embodichain/lab/gym/envs/base_env.py) 先恢复物理对象再调用 dataset 保存；[EmbodiedEnv.reset](../../embodichain/lab/gym/envs/embodied_env.py) 清空活动任务桥并播种首帧 | 旧 episode 验收证据必须先冻结；初态准备与观测重播种需要独立生命周期入口 | +| [AsyncLeRobotRecorder](../../embodichain/lab/gym/envs/managers/async_datasets.py) 队列无界，没有逐 commit 回执 | 同步提交协议先行；之后补队列反压、完成回执和错误传播 | +| 本机 LeRobot 0.4.4 的 `save_episode()` 写数据后仍保持 Parquet writer 打开，`finalize()` 才完成 footer 等收尾；项目 recorder finalize 后不能继续复用 | 提交确认需要明确封存屏障与读回检查；不能将 episode 写入函数返回作为已持久化 | +| CLI 实际注册位于 [cli/main.py](../../embodichain/cli/main.py)，`__main__.py` 仅转发 | 新启动器在现有 dispatcher 注册，保持根命令的懒加载 | + +## 3. 模块与契约归属 + +按模块归属 review 采用以下布局；未标为已实现的生成宿主部分仍为后续实施计划，只有职责变复杂时再拆文件： + +```text +embodichain/lab/sim/motion/ + __init__.py # 按需暴露四个子包，不汇总其类和函数 + solvers/ # FK、IK、微分运动学 + planners/ # 路径、碰撞约束、时间参数化 + workspace/ # 离线可达性分析、缓存与运行时采样 + +embodichain/lab/sim/motion/trajectory_augmentation/ + contracts.py # case/snapshot/template/candidate/episode/receipt 与端口协议 + cfg.py # 两类 @configclass；只持配置值及注册 ID + operators.py # 因子提议、自由段几何和时间变体 + coverage.py # 几何家族、去重、覆盖计算 + session.py # case 分区、预算、候选池、配额预留、结果归账 + +embodichain/lab/trajectory_generation/ + initial_state.py # 已实现：可信 profile、全批 host、PreparedBatch epoch + runner.py # 已实现：手写 qpos/free-motion 同步编排与 generation_report + execution.py # 已实现：sim/Gym qpos 执行、实际标签与因果证据冻结 + config.py # generation.yaml 严格加载、可信注册与能力检查 + integrations/ + planning.py # 已实现 free/no-held EEF/qpos 的 env_rows 适配与采样检查 + handwritten.py # 显式模板与 DemoSegment 适配 + atomic.py # 已实现：离线 PickUp 阶段模板；运行时 ActionPlan 适配另行落地 + contact.py # 已实现：PickUp 完整几何和 CPU 物理子步接触验收 + sim.py # 已实现物理初态适配；rollout 控制/计时/观测归 execution.py + gym.py # Gym 初态、正常 env.step、demo/Task Program 桥 + validation.py # 注入路径检查与实测任务/质量验证 + sinks.py # 已实现同步 LeRobotEpisodeSink，已接入真实自由运动采集 + +embodichain/lab/scripts/generate_trajectories.py # 拟新增启动器 +``` + +接口 review 重点： + +1. `SceneCase / initial_state_id`、`candidate / family / attempt`、`slot / runtime_epoch` 三组身份独立；arena 平移通过坐标映射处理，不能进入候选随机种子。 +2. 快照及候选不持有活环境或可变 backend；明确张量所有权，避免表面只读而底层数据仍被宿主更新。 +3. `CandidateTrajectoryBatch` 使用 `(C,N,D_full)`、逐行有效长度和显式时间。固定首点时间、到达间隔、padding 与阶段事件区间语义；工具事件、阶段、失败状态与候选索引一同映射。 +4. Session 只接收快照与结果；Runner 管生命周期；sim/Gym 各有一个步进所有者。Gym 的控制周期取 `env.step_dt`。 +5. 原子候选执行建议采用受限的“初始计划供应端口”：绑定新的 request/context 后消费指定候选的完整初始计划；后续 recovery 保留既有机制。具体签名在 PR 1 固定，不能通过写私有 session 字段实现。 +6. `EpisodeSink` 的提交输入有稳定 episode/commit ID；回执携带原候选身份。定义何时主数据、必需媒体及谱系均完成持久化，enqueue 或单个 writer 函数返回均不自动算 committed。 +7. source、validator、prepare、tolerance、motion-limit profile 使用可信注册 ID；YAML 不执行任意 dotted import。模板和 profile 的约束只能被 job 收窄。 + +能力声明分别描述多 seed、all-solutions、关节分支保持、精确路径验证、独立候选 batch、初态复制和持久化屏障，不能用一个通用 `supports_batch` 布尔值替代。M1 先支持明确的 qpos 控制/标签 profile 和工具事件；其他表示需先实现并验证转换。 + +## 4. 实施顺序与 PR 划分 + +以下编号用于 review 和依赖规划，不表示本次创建 PR。每个 PR 随实现补对应测试与公共 API 文档。 + +### M0：建立可执行基准与能力清单 + +**交付内容** + +- 固定一个外部给定的 cube case、机器人初态、工具标定、控制周期和传感器设置。 +- 定义真实成功判据：物体按任务要求抬升并保持、抓持关系有效、无掉落；容差、持续时间、跟踪误差和质量上限均有数值及单位。 +- 验收未扩增的手写 EEF 基准与原子 PickUp。记录路径长度和实际时长，建立同起点、同任务/模式的 ratio 比较基准。 +- 检查路径碰撞、阶段接触许可、持物几何、初态恢复和 writer 持久化的实际能力。缺项列为 M1 必须补齐的工作。 +- 加入基础计时，保存规划、reset/settling、物理、渲染、写盘时间，后续优化沿用同一对照条件。 + +**退出条件**：基准能重复运行；未支持的必需验证有明确实现方案。不能用关闭检查的方式进入专家数据采集。 + +### PR 1：公共契约、配置和最小 Session + +**依赖**：M0 的任务与 profile 定义。 + +- 新增公共值对象、端口、两类配置及严格解码；固定状态/时间/动作表示、候选消费和提交确认协议。 +- 实现 case 分区、稳定 ID、本地 RNG、规划/rollout/时间预算及基本有界候选池。状态区分 `proposed / planned_valid / rollout_attempted / validated_accepted / pending_write / committed`。 +- 实现最小的家族归属、去重与配额预留；同一路径的 timing 变体共用 geometry family，不增加空间覆盖。复杂价值调度留到 PR 9。恢复初态不清空 job 历史。 + +**验收**：纯 CPU/张量测试；算法无直接 Gym/backend 依赖，motion 父包按需加载子包；slot/chunk 改变不改变纯因子提议；时间变体不刷空间覆盖；重复回执、失败及尾批不会重复计数或超额；未知/无能力配置拒绝启动。 + +### PR 2：sim/Gym 初态准备与恢复生命周期 + +**依赖**：PR 1。 + +**当前落地**：`SimInitialStateAdapter` 严格支持一个配置所有的固定基座 robot 及完整批次普通刚体，保存关节/根/物体状态、速度、drive target 与关节 effort，并在任何写入前检查完整 schema、结构签名、SE(3) 和关节限位。`FixedSceneHost` 以调用者提供的各行 case 建立初态，恢复前使旧 binding 失效，经过可信 profile 准备、物理恢复、settling 和验证才发布新 epoch。profile 还必须覆盖物理快照之外的控制历史、任务状态与固定物理/视觉/传感器条件。 + +Gym 接口增加全批 generation lease，租约内拒绝普通 reset 并抑制自动 reset；Gym 和纯 sim host 登记同一 simulator batch 所有权，不允许两个不同 owner 并存。`prepare_generation_episode()` 先丢弃旧录制，再依次准备、恢复、settle、清空 standard episode manager 历史、verify，确保 validator 看到新 episode 的 manager 状态；验证通过后才获取首观测并播种记录。失败时禁止继续 step。调用者必须在进入该生命周期前取得旧 episode 的独立证据。该接口不自动保存数据，不调用 reset event 或重播种环境 RNG，也不声明任意 controller 已被默认覆盖。 + +物理根位姿 setter 可能推进整个 world，因此恢复先写根，再恢复所有其余状态；异常不会伪装成成功。已通过物理关闭/开启两种模式的真实 CPU backend 状态恢复 smoke；该测试与 mock 生命周期验证均不替代参考任务在接触和长时间 settling 下重复恢复的验收。 + +- 实现 `acquire_case / restore_initial / verify_initial` 与 slot epoch。先完成单槽，再覆盖全批中每行不同 case。 +- 启动检查要求 Runner 独占整个被 reset 的宿主 batch，并在采集期间关闭隐式 auto-reset；不能让外部活动行被本 job 的恢复连带重置。 +- 明确恢复 robot/rigid-object 状态和速度、夹爪与控制目标、任务/manager 初态；经过 settling 后验证。首版不支持的实体或初始约束明确报错。 +- Gym 增加受控准备入口：旧证据冻结与提交载荷取得 → episode 清理 → 所需确定性初始化/状态恢复 → settling → standard episode manager reset → 验证 → 刷新观测 → 播种首帧 → 新 binding。 +- 显式区分改变布局的随机事件与必需初始化。inner restore 不消耗外层布局 RNG；运行中改变 case 固定条件的事件须受宿主策略约束。 + +**验收**：多次恢复后物体/关节/速度/控制状态符合 profile；观测首帧对应恢复后的现场；准备段不写训练帧；失配和旧 epoch 拒绝执行；全批 reset 不影响其他独立宿主。 + +### PR 3：episode 冻结、同步 sink 与提交确认 + +**依赖**：PR 1；与 PR 2 可并行开发。 + +**当前落地**：`LeRobotEpisodeSink` 接收已冻结且验收通过的 `ExpertEpisode`，每条独立封存一个本地 LeRobot dataset shard。主数据保留 T 个因果训练帧，`terminal.npz` 保留末次观测和全部 T+1 个实测时间，`episode.json` 保存谱系、动作表示、验证、阶段和 metadata；collection manifest 才是已提交片的目录。同步写入完成后重新打开并读取每帧/图像、全部必需 sidecar 与 manifest，成功后返回 confirmed receipt。 + +首版只接受固定整数 fps、支持的数值观测向量与 uint8 RGB、float32/float64 动作；实测时间必须在绝对容差内符合宿主固定时钟。单条 payload 有 raw tensor + metadata 字节上限。输入错误在 episode 写入前抛出；持久化失败返回原身份的失败回执。相同 commit ID 必须携带相同 payload：已 confirmed 时只读核验，不重写产物；尚未确认时可复用已封存可读片，未提交且不完整的片在原路径重建，不新增逻辑 episode。数值读回检查 Parquet 原始精度，RGB 通过 LeRobot 解码。`drain()` 没有延迟回执，`close()` 不隐式追加保存。仅支持新/空输出目录和当前进程的单写入者，不声明重启恢复、跨文件事务或断电持久性。 + +`QposRolloutExecutor` 已冻结真实 rollout 的 T+1 观测、T 个实际 qpos 命令和实测时间；`GenerationRunner` 完成后续实测验收、Session 归账和 sink 提交。sink 只消费调用者提供的验证，不重新计算碰撞或任务成功。真实 UR5 正例已完成整个写入/封存/读回链路。 + +- 定义统一 `ExpertEpisode`、候选谱系和实际动作转换 profile；复用已有因果配对与有效帧标记。 +- Gym 使用 DatasetManager/LeRobotRecorder 适配；sim 使用宿主无关 sink。优先抽取可复用的序列化/写入部分，避免模拟 Gym env 对象。 +- 补逐 episode 的提交结果。主记录、媒体和 metadata/sidecar 部分成功时能识别已写结果，重复提交同 commit ID 不生成第二条逻辑 episode。 +- 回执按提交身份归账；只有持久化确认后更新正式覆盖。Runner 明确调用 drain/finalize 并传播错误。 +- 建议 LeRobot 从同步版本就采用有界分片：达到 pending 条数/字节上限或本轮剩余配额后停止接收，封存当前 recorder，读回主数据/媒体/sidecar，再发回执；需要继续时创建下一分片。job manifest 保存 commit 到分片/episode 的映射，封存期间 pending reservation 不释放。收尾失败时保留未完成状态并回传错误，不重复写主记录。 +- 首版持久化语义固定为所有必需产物已关闭、可重读并完成错误检查；进程崩溃恢复、断电持久性和跨文件事务另行定义，不暗含在普通回执中。 + +**验收**：reset 后仍能保存旧 episode 的冻结证据及末次 observation；拒绝行不写专家集；覆盖只取已提交的实际轨迹；注入 writer/sidecar 失败后不误计数；重复提交与迟到回执可正确处理；真实 LeRobot 完成一次写入→封存→重新打开并读取的验证。大于单片容量的 job 能持续推进，pending 满额不死锁。 + +### PR 4：候选规划、时间与验证的最小公共路径 + +**依赖**：PR 1；与 PR 2/3 可并行开发。 + +**当前落地**：`EEFPath` 保存 local arena-frame TCP 的显式样本、初始锚点、arrival intervals、来源行及阶段；可提供 control-part 顺序的 solved qpos，保持分支并经 FK 验收，不重复 IK。`EnvRowMotionPlanner` 把 C 个候选按真实 B 行分轮，每轮同一来源行至多处理一个候选；返回原序且来源身份对齐的 full-joint 候选与逐候选验证。自由 qpos 路径按 joint step 上限在相邻样本和阶段边界间加密，交给支持精确关节样本验证的 backend 检查。 + +当前只支持完整 `free` 标注、显式无持物，且全部未受控关节保持 snapshot 与 `robot.cfg.init_qpos` 的初值。contact/hold/未标注路径、持物、夹爪等锁定关节变化、无 backend 能力或采样超过容量返回 `unavailable`；碰撞和非法路径返回 `failed`。结果仅提供 `path_collision` 证据；实际速度/加速度、任务成功、接触与物理质量仍是独立 gate，不把离散加密解释为连续碰撞证明。 + +- 实现显式 `env_rows`：任意 C 按真实 B 分轮，一轮每个物理来源行至多处理一个候选，其余安全占位；维护候选到 snapshot/root/world 的映射。 +- 手写 EEF 走注入的 IK/规划端口；支持固定已解关节目标及分支连续性约束。未标注 qpos 默认仅原样执行；未受控关节保持原值。 +- 首版每个自由段只用一种 via-point 算子；retime 后按宿主周期重新采样，重算事件索引、速度/加速度与阶段连接。 +- 实现所选任务必需的路径检查：自碰撞、环境、段间路径、阶段接触许可、持物扫掠；复用 `MotionGenerator.validate_joint_trajectory()`，补足其不覆盖的部分。 +- 记录 `not_run / passed / failed / unavailable` 和连续指标；必需项 unavailable 的候选只能进入规划库。 + +**验收**：C≠B 映射往返、空输入、失败 seed、NaN/Inf、padding 隔离、固定端点、实际 IK 分支、变长时间和事件顺序均正确；故意碰撞或携物扫掠碰撞被拒绝。 + +### PR 5：共享 Runner 与手写来源端到端闭环 + +**依赖**：PR 2、3、4。 + +**当前状态：handwritten qpos/free-motion 同步闭环已落地。** `GenerationRunner` 与 `MotionLimitsProfile` 要求显式传入现有 host、planner、executor、sink 和可信限速配置，对齐 source/template/profile/validator/tolerance ID 与 control_dt/fps。运行入口接收每个物理行的固定 case 和 qpos reference，提议允许的残差/retime，规划检查后先预留配额与容量，再执行。`QposRolloutExecutor` 采集实际命令与实测观测/时间，冻结证据后由 Runner 重验实际碰撞、速度/加速度、路径/时长质量和任务成功，再提交。`generation_report.json` 记录配置、包版本、计数、audit 与目标达成状态；`last_failures` 将无完整 transition 等失败原因有界保留到 audit。Runner 为 single-use，并在退出时关闭传入的 host/sink。 + +自由运动路径使用原 `EnvRowMotionPlanner` 的锁定关节约束。新增 PickUp 路径必须将同一个 `PickUpMotionValidator` 传入 Runner 和 executor，才能允许目标物体与 mimic 运动并采集物理子步证据。它接收离线原子模板；完整 Atomic Runtime、Gym PickUp、M1 和统一 CLI 仍未完成。 + +真实正例使用纯 arm UR5、普通重力、CPU physics + CUDA cuRobo、B=1：1 秒内 21 个观测/20 个命令，实际位移约 `0.078021 rad`、终点误差 `0.003212 rad`、最大跟踪误差 `0.010872 rad`，全部 gate 通过且 LeRobot 读回 confirmed，`committed=1`。Panda 默认重力负例中 task/quality/dynamics 通过，但最大手指漂移约 `3.7853e-5` 超过 locked model 的 `1e-6` 比较阈值，实测路径验收 unavailable，`committed=0`。候选不发出夹爪变化命令不能代替实测锁定关节一致性。两项真实测试均通过；未据此验收真实 Gym 采集、多行或连续多 episode。 + +在仓库根目录、已安装仿真/CUDA/cuRobo/LeRobot 依赖时运行： + +```bash +python examples/sim/motion/trajectory_generation/free_motion.py --output /tmp/ur5-free-motion +pytest tests/lab/trajectory_generation/test_runner_real.py --run-gpu -m gpu -q +``` + +输出目录必须新建或为空；示例显式构造与隐式地面匹配的 `ground_proxy`，全部普通刚体均动态更新碰撞 pose。成功输出 `generation_report.json`、已提交 `manifest.json` 和每条 episode 的 LeRobot dataset/`terminal.npz`/`episode.json`。`--robot panda` 可在另一个空目录运行上述拒绝案例；未达目标返回 exit 1。这是具体示例脚本，统一 generation CLI 仍待 PR 6。 + +Runner 要求全部普通刚体的物理 UID 同时存在于 backend 的 collision-world 与 dynamic-pose ID 集合，每轮以捕获的各行初态更新；这里 dynamic 表示 backend 可更新 pose，不要求物理刚体为动态。仅烘焙静态 world 无法保证准备后姿态或多行不同 case 一致；B>1 且相对布局不同时需 per-environment backend world。语义 aliases 当前不接入该 Runner，隐式平面等未进入刚体注册表的几何由可信 profile 认证碰撞覆盖。executor 的逐 episode 字节上限必须不超过 sink 与 pending 上限。 + +- 在 demo executor 增加显式候选 `segments` 输入，与旧 demo hook 路径做明确互斥/选择；支持 iterator 按需消费，不提前耗尽含反馈的旧生成器。 +- Runner 编排规划、准备、执行、验收、提交；实现 `full_batch`，完成行安全 hold，停止生成训练帧。纯规划失败不 reset。 +- 手写 × Gym 使用正常 `env.step()`；控制器目标走 `ControllerAction`，不重复 action pre 处理。手写 × sim 独立完成控制、整数 physics substeps、传感器和记录。 +- 候选结束前采集任务/质量证据；需要持续成立的效果在批次收尾时再检查。终态证据冻结之前不得触发自动或外部 reset。超时和取消关闭新工作并收尾在途数据。 + +**验收**：两条手写路径执行的是指定候选；记录满足 `observation_t → action_t → observation_{t+1}`；无重复步进、旧首帧、padding/hold 污染；即使 done/truncated 也先冻结再 reset；无法取得全批生命周期所有权时拒绝启动；连续采集多条已确认 episode。 + +### PR 6：原子 PickUp 接入、启动器与 M1 发布验收 + +**依赖**：PR 4；端到端合并验收依赖 PR 5。 + +- 原子适配器导出 PickUp 阶段模板，首版固定合法 contact、保留一条有效分支；候选绑定到当前 case/slot 后才构造新的请求和 ActionPlan。 +- 从 PickUp 抽出候选生成与计划物化的共享逻辑；空 grasp、padding 与失败 seed 隔离在首批实现中修复，不复制另一套 PickUp 算法到扩增器。 +- 为现有原子执行增加正式的候选消费端口；sim 复用 ExecutionSession/Runner，Gym 经 Task Program 命令桥下发，保持效果验证和命令语义。 +- Gym 保留桥的正常耗尽、post-policy、validator 和 abort handshake；命令数组执行完不能独立替代任务验收。 +- M1 不采集恢复专家数据：若既有 engine recovery 被触发，沿原协议安全收尾，只保留 audit 并拒绝该 episode 进入专家集。专用恢复数据留待具备边界隔离和标签协议后扩展；在线默认逐环境选一个 winner 的调用保持原行为。 +- 新增独立启动器，建议命令名 `generate-trajectories`，显式加载宿主配置与 `generation.yaml`；注册到 `cli/main.py`，产出已解析配置、版本、audit 与最终计数。 +- 提供 M1 四组合 fixture 与可运行示例；同步公共 API 文档和受影响项目上下文。 + +**M1 验收门槛**：四组合都用无 GUI 真实物理运行；建议每组合提交 8 条合格 episode，另覆盖规划失败、物理失败、恢复失配、写入失败和尾批场景。8 条仅为工程烟测规模，不作为数据效果结论。不得以预算耗尽或命令完成替代成功标准。 + +### PR 7：同 case 副本池 + +**依赖**：M1。 + +- 实现 `grouped_replicas`、副本容量检查、case 分配及跨 arena 坐标变换；只在整批 episode 收尾后重绑定。 +- 同 case 副本共享不变输入，物理状态和运行时 binding 独立；每轮复制/恢复后重新校验初态。 + +**验收**：至少覆盖 `1 case × 4 slots` 与 `2 cases × 2 slots`；同候选换兼容副本后映射正确；不同 case 的 world、覆盖和数据不串用。 + +### PR 8：逻辑候选批量适配与稳定容量桶 + +**依赖**：PR 4、M1；可与 PR 7 并行开发。 + +- 新增 `generate_candidates()` facade，以能力声明区分 `candidate_buckets` 与 `env_rows`。 +- 支持候选来源重复行、显式 robot-root/初态映射、shared/per-candidate collision world 和安全 padding。复用现有 cuRobo backend 缓存，不伪造 `num_instances`。 +- 首先适配一个明确支持的 backend；其他 backend 保留显式 `env_rows` 路径。限制桶数量、warmup 与显存占用。 + +**验收**：例如 C=11、B=4、桶容量 16 的索引与失败 mask 正确；改变桶不能串世界或丢候选；与 M1 同预算比较规划成本及端到端收益。 + +### PR 9：依赖缓存、覆盖驱动选择与时长分桶 + +**依赖**：M1；整体验收包含 PR 7/8。 + +- 缓存键纳入 source revision、case、任务/标定/起点、阶段依赖和执行限制；只缓存可复用值，不缓存旧 session/ActionPlan。 +- 分离几何与时间家族，按阶段进度做近重复判断；正式覆盖由实际且 committed 的轨迹更新，在途预留失败后释放。 +- 实现低/高水位补池、按需分层展开、时长桶含溢出桶、价值/剩余成本排序、探索比例与每 case 最低配额。 +- 加入独立探测、预算和饱和停止条件；所有反馈按稳定逻辑顺序合并。 + +**验收**：纯 timing 变化不刷空间覆盖；case/起点/标定变化正确失效；旧回执不归入新 slot；缓存和池满足条数/字节上限；尾批不超额。 + +### PR 10：有界异步写入与流水线 + +**依赖**：PR 3、9。 + +- Async recorder 接入相同 commit 协议；enqueue 前取得独立 payload，按条数和字节限流。 +- writer 阻塞向候选生产传播反压；单请求超限明确拒绝或走事先声明的流式 sink。 +- 先重叠 CPU 整理/写盘；planner 可变 world 与求解受串行保护。GPU 规划和物理/渲染默认不重叠，只有实测收益及正确性通过后才启用。 + +**M2 验收门槛**：同任务、相同 case/槽数/质量门槛/每 case 数据目标下，逐项消融 PR 7–10；报告 committed/min、覆盖增益/min、到达目标时间和内存峰值。没有实测收益的优化不设为默认。 + +### M3：扩大运动覆盖，按来源/技能拆 PR + +**依赖**：M1 的验证和候选协议；批量路径复用 M2。 + +1. **多 grasp / 多 IK**：graspkit 增加 masked batch 辅助协议并保留旧 ragged 接口；解析多解/多 seed 构型去重，逐分支传播成功 seed；PickUp 推迟 winner，保持实际关节分支到 lift。验证空 grasp、失败 seed 隔离和最终分支去重。 +2. **手写多段与 PickUp → Place**：接入父前缀/持物变换及后缀可行性；从实际执行边界状态生成续段；中间阶段不 reset。完整任务与成功片段独立标识。 +3. **逐技能推广**：AxisAlign、Slide、OpenDoor、HandOver 分别补特定约束与真实验收；HandOver 单独覆盖双臂联合碰撞和交接时序，不合并为一个大 PR。 +4. **数据使用与学习验收**：基于 job manifest 汇集已封存分片，按 trajectory family/parent 分组 train/validation，限制时间近重复和公共前缀权重；同数据预算比较模仿学习成功率，保留未见起点/运动参数组合。候选 ID、seed、未来成功标记只作 audit。 + +**M3 验收门槛**：增加的是通过物理/质量验收的运动覆盖；学习收益以实测报告为准。固定场景结果不宣称布局泛化。 + +### M4：按实测瓶颈选择后续工作 + +逐行补位必须先证明 reset、速度清理、控制器、任务桥、观测历史、记录和碰撞绑定全链路行隔离；所有 settling 仍由唯一 batched step 推进。checkpoint 分叉必须另行建立完整状态恢复与前后缀连续性协议。真实 recovery 数据需要可达前缀和可靠后缀专家。这些能力均不作为 M1–M3 的交付依赖。 + +## 5. 依赖关系与并行组织 + +```mermaid +flowchart LR + A[M0 基准与能力] --> B[PR 1 契约与 Session] + B --> C[PR 2 初态恢复] + B --> D[PR 3 提交协议] + B --> E[PR 4 规划与验证] + C --> F[PR 5 手写闭环] + D --> F + E --> F + E --> G[PR 6 原子适配] + F --> H[M1 四组合验收与启动器] + G --> H + H --> I[PR 7 副本池] + H --> J[PR 8 候选批量] + I --> K[PR 9 缓存与调度] + J --> K + K --> L[PR 10 有界流水线 / M2] + H --> M[M3 多分支与多阶段] + L -. 推荐先完成效率验证并复用其能力 .-> M +``` + +PR 1 合入后可并行推进宿主生命周期、sink 和规划验证三条线。PR 6 的来源适配可与 PR 5 并行，最终由同一四组合矩阵验收。排期应以 M0 暴露的碰撞/持物验证和持久化缺口为依据，暂不对未实测工作承诺固定工期。 + +## 6. 验证与交付物 + +| 验证层 | 建议落点及重点 | +|---|---| +| 纯张量核心 | 新增 `tests/sim/motion/trajectory_augmentation/`：schema、身份、RNG、去重、覆盖预留、预算和队列；独立检查 motion 懒加载与算法直接依赖边界 | +| 初态准备 | [物理状态适配](../../tests/lab/trajectory_generation/test_sim_initial_state.py) 与 [宿主生命周期](../../tests/lab/trajectory_generation/test_initial_state_host.py)：独立快照、全批预检、固定条件漂移、恢复误差、失败和旧 epoch 拒绝 | +| 规划/原子回归 | 扩展 [atomic_actions 测试](../../tests/sim/atomic_actions/)、[MotionGenerator](../../tests/sim/motion/planners/test_motion_generator_batched.py)、[BasePlanner](../../tests/sim/motion/planners/test_base_planner.py)、[graspkit](../../tests/toolkits/test_grasp_pose_generator.py) | +| 宿主/记录回归 | 扩展 [Demo](../../tests/gym/envs/test_demo.py)、[trajectory_state](../../tests/gym/utils/test_trajectory_state.py)、[DatasetManager](../../tests/gym/envs/managers/test_dataset_manager.py)、[同步 recorder](../../tests/gym/envs/managers/test_dataset_functors.py)、[异步 recorder](../../tests/gym/envs/managers/test_async_dataset_functors.py) | +| 集成与失败注入 | 扩展 `tests/lab/trajectory_generation/`：fake-host 契约测试、四组合真实仿真、唯一 step、恢复失配、终态冻结、延迟/重复/失败回执 | +| CLI 与兼容 | 扩展 [tests/test_main.py](../../tests/test_main.py)；旧 run-env/demo 调用、逐环境单 winner、唯一 env IDs 和 ActionPlan 效果协议保持有效 | +| 性能与数据效果 | 按项目 benchmark 规范新增 generation benchmark；M0 保留基线，每个效率 PR 附同条件对照，M3 再增加学习评估 | + +每个里程碑提交：已解析配置与能力清单、代码/backend 版本、固定 case 与 profile、候选/episode/commit 谱系、失败统计、计时与内存报告、可读的数据样本和复现实验命令。公共导出变动同时补 API 文档；按改动范围更新 env-framework、manager-functor、atomic-actions、motion-planning 上下文。 + +实施时按项目技能完成测试与 PR 检查：`add-test`、`pre-commit-check`；涉及公共 API 用 `update-api-docs`，新增 benchmark 用 `benchmark`。每个 PR 做比例适当的验证；四条真实路径和效率对照属于里程碑发布门槛，不能仅由 mock 测试替代。 + +## 7. 实施约束与后续验收重点 + +1. M1 以 **UR5 单物体 PickUp + 手写 EEF** 作为共同 fixture，要求四组合全部跑通；当前纯 arm 自由运动示例只验收其支持的子集。 +2. 按 **先 env_rows/full_batch/同步写入，再副本池与候选容量桶** 的顺序交付。 +3. **初态恢复、指定候选消费、阶段碰撞/持物验证、逐提交确认** 都是 M1 前置条件，不按吞吐优化项后移。 +4. 候选通过原子执行的受限计划供应端口接入，保持原有执行、效果验证和 recovery 所有权。 +5. M0 实测后再固定成功容差、碰撞 backend 和工期；吞吐和学习增益按同预算实验验收。 diff --git a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst index 79d1b3614..bc8269a21 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst @@ -56,6 +56,36 @@ Environment Classes :members: :exclude-members: __init__, class_type +Generation Preparation +---------------------- + +``BaseEnv.acquire_generation_lease(owner)`` reserves the complete environment +batch, suppresses automatic resets, and rejects ordinary resets until the same +owner calls ``release_generation_lease``. ``generation_epoch`` invalidates old +runtime bindings when ownership or episode preparation changes. A lease is an +ownership contract; callers must still serialize host access. Gym leases and +direct simulation generation hosts register the same simulator-batch ownership, +so they cannot reserve that batch for different owners. + +``EmbodiedEnv.prepare_generation_episode`` requires the active lease and four +explicit preparation, restoration, settling, and verification callbacks. It +discards buffered recordings, so the caller must freeze the old episode first. +After settling it resets episode managers before verification, so validation +observes the new episode's state. It obtains the initial observation and seeds +recording history only after verification succeeds. Failed preparation leaves +stepping disabled; normal reset events and RNG reseeding are not part of this path. + +The shared +:class:`~embodichain.lab.trajectory_generation.initial_state.FixedSceneHost` +assembles this lifecycle with physical state restoration and a trusted task +profile. See :doc:`/overview/trajectory_generation` for usage and limitations. + +``BaseEnv.observe_generation_commands(owner, observer)`` installs one serialized +callback for the active lease. It runs after controller submission and before +physics, allowing the collector to copy actual controller targets. An exception +occurs after the command was submitted, so the attempt still counts and the +collector must stop safely. + Controller-ready Actions ------------------------ @@ -77,6 +107,13 @@ containing one or more semantic subtasks. Segment action iterables may be lazy, and the common executor records per-environment lengths, terminal status, and segment spans. +``execute_demo_episode`` accepts ``step_observer`` to consume the existing +``env.step`` result and its active-row mask before terminal handling, without +querying observations again. With explicit segments, ``row_step_limits`` provides +one command count per row; zero skips a row and finished rows receive holds while +remaining rows finish. Their post-completion holds are not recorded as training +frames. Omitting these arguments preserves ordinary segment execution. + .. currentmodule:: embodichain.lab.gym.envs.demo .. autoclass:: DemoExecutionCfg diff --git a/docs/source/api_reference/embodichain/embodichain.lab.gym.rst b/docs/source/api_reference/embodichain/embodichain.lab.gym.rst index 3fefee099..0a48c04ec 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.gym.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.gym.rst @@ -39,6 +39,7 @@ Base Environment Classes :members: :inherited-members: :show-inheritance: + :exclude-members: acquire_generation_lease, release_generation_lease, generation_epoch, observe_generation_commands The foundational environment class that provides the core functionality for all EmbodiChain RL environments. This class extends the Gymnasium ``Env`` interface with multi-environment support and robotic-specific features. @@ -56,6 +57,7 @@ Embodied Environment Classes :members: :inherited-members: :show-inheritance: + :exclude-members: acquire_generation_lease, release_generation_lease, generation_epoch, prepare_generation_episode, observe_generation_commands An advanced environment class that provides additional features for embodied AI research, including sophisticated observation management, event handling, and multi-modal sensor integration. @@ -67,6 +69,10 @@ Embodied Environment Classes Configuration class for embodied environments with extended settings for lighting, observation management, and advanced simulation features. +Generation lease and preparation methods are documented with the owning +:doc:`environment APIs ` and the +:doc:`fixed-scene generation guide `. + Utilities Module (utils) ------------------------- @@ -134,4 +140,3 @@ Miscellaneous Utilities .. automodule:: embodichain.lab.gym.utils.misc Miscellaneous utility functions for environment development and debugging. - diff --git a/docs/source/api_reference/embodichain/embodichain.lab.rst b/docs/source/api_reference/embodichain/embodichain.lab.rst index 43cd72639..2e40118d9 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.rst @@ -10,7 +10,8 @@ The ``lab`` package is EmbodiChain's robotics laboratory. It owns the provider-independent Task Program language (including its Semantic Call contracts), the simulation core (``sim``), the Gymnasium-compatible environment framework (``gym``), real-device controllers (``devices``), and -browser visualization (``visualization``). +browser visualization (``visualization``), and fixed-scene generation host +integration (``trajectory_generation``). .. rubric:: Submodules @@ -21,6 +22,15 @@ browser visualization (``visualization``). gym sim visualization + trajectory_generation + +Trajectory Generation +--------------------- + +.. toctree:: + :maxdepth: 1 + + embodichain.lab.trajectory_generation Browser Visualization --------------------- diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.planners.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst similarity index 88% rename from docs/source/api_reference/embodichain/embodichain.lab.sim.planners.rst rename to docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst index 3af9a9efc..1b4dbd01d 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.planners.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst @@ -1,7 +1,7 @@ -embodichain.lab.sim.planners -========================================== +embodichain.lab.sim.motion.planners +=================================== -.. automodule:: embodichain.lab.sim.planners +.. automodule:: embodichain.lab.sim.motion.planners .. rubric:: Classes @@ -18,7 +18,7 @@ embodichain.lab.sim.planners PlanResult PlanState -.. currentmodule:: embodichain.lab.sim.planners +.. currentmodule:: embodichain.lab.sim.motion.planners Base Planner ------------ diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.rst new file mode 100644 index 000000000..d4d8d15c6 --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.rst @@ -0,0 +1,45 @@ +embodichain.lab.sim.motion +========================== + +.. automodule:: embodichain.lab.sim.motion + +The ``motion`` package groups the robot motion capabilities used by simulation +objects, atomic actions, and environment integrations. Import the required +capability from its owning subpackage: + +.. code-block:: python + + from embodichain.lab.sim.motion.solvers import SolverCfg + from embodichain.lab.sim.motion.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.motion.workspace import RobotWorkspaceCfg + from embodichain.lab.sim.motion.trajectory_augmentation import TrajectoryTemplate + +The parent package exposes these four subpackages through lazy attributes; +it does not eagerly import them or re-export their classes and functions. +This keeps their initialization boundaries separate: robot configuration can +use solver and workspace contracts without ``motion`` also initializing the +planner layer. Normal ``lab`` and ``sim`` package initialization still applies. + +Solvers own FK, IK, and differential kinematics; planners turn targets into +paths and timed trajectories; workspace provides offline reachability analysis +and runtime cache queries; trajectory augmentation owns candidate variation, +coverage, and generation accounting. Physical rollout, reset, task validation, +and dataset persistence belong to explicit host integrations. + +.. currentmodule:: embodichain.lab.sim.motion + +.. autosummary:: + :nosignatures: + + solvers + planners + workspace + trajectory_augmentation + +.. toctree:: + :maxdepth: 1 + + embodichain.lab.sim.motion.solvers + embodichain.lab.sim.motion.planners + embodichain.lab.sim.motion.workspace + embodichain.lab.sim.motion.trajectory_augmentation diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.solvers.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst similarity index 94% rename from docs/source/api_reference/embodichain/embodichain.lab.sim.solvers.rst rename to docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst index f7eeab076..be67db31a 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.solvers.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst @@ -1,7 +1,7 @@ -embodichain.lab.sim.solvers -========================================== +embodichain.lab.sim.motion.solvers +================================== -.. automodule:: embodichain.lab.sim.solvers +.. automodule:: embodichain.lab.sim.motion.solvers Overview -------- @@ -40,7 +40,7 @@ numerical (``Pinocchio``, ``Pink`` with null-space posture tasks, NeuralIKSolverCfg NeuralIKSolver -.. currentmodule:: embodichain.lab.sim.solvers +.. currentmodule:: embodichain.lab.sim.motion.solvers Base Solver ----------- diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.trajectory_augmentation.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.trajectory_augmentation.rst new file mode 100644 index 000000000..2cdce3eab --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.trajectory_augmentation.rst @@ -0,0 +1,246 @@ +embodichain.lab.sim.motion.trajectory_augmentation +================================================== + +The :mod:`embodichain.lab.sim.motion.trajectory_augmentation` package provides +qpos contracts, constrained trajectory operators, measured coverage, and a +generation session for simulation expert trajectories. The core algorithms do +not directly import Gym or own simulation stepping. Public imports pass through +``embodichain.lab`` and ``embodichain.lab.sim`` initialization and therefore +require the normal simulation dependencies. + +This is the motion core for fixed-scene expert generation. Full-batch physical +initial-state restoration, supported free-motion EEF/qpos checks, actual qpos +execution, the synchronous runner, and episode persistence are provided by the +separate :doc:`generation host API `. +That layer supplies actual observations and commands, validation evidence, and +persistence confirmations. The low-level candidate execution entry points in the +simulation and Gym layers are separate from this value-only API. + +.. currentmodule:: embodichain.lab.sim.motion.trajectory_augmentation + +.. autosummary:: + :nosignatures: + + CandidateIdentity + CandidateTrajectoryBatch + CommitReceipt + ExpertEpisode + MotionSnapshot + SceneCase + TrajectoryAugmentationCfg + TrajectoryGenerationJobCfg + TrajectoryPhase + TrajectoryTemplate + ValidationCheck + ValidationResult + TrajectoryDescriptor + describe_trajectory + CoverageIndex + rotate_grasp_about_object_axis + joint_residual + retime + validate_motion_limits + GenerationSession + +Values and Evidence +~~~~~~~~~~~~~~~~~~~ + +``SceneCase`` separates scene and initial-state identity from physical execution +slots. ``MotionSnapshot`` copies robot positions and velocities, the robot root +pose, entity poses, and dependency revisions. Pose validation checks finite +proper SE(3) transforms, including rotation orthogonality and handedness. +Tensor inputs are detached and cloned at construction; consumers must still +treat the resulting owned tensors as read-only. + +Templates declare the complete joint order and use explicit qpos values. +``dt[0]`` is zero and each subsequent ``dt`` is the positive arrival interval +from the preceding sample. Phase ranges are half-open sample intervals +``[start_index, stop_index)``. Unannotated paths permit replay only; operators +also require explicit template, phase, and controlled-joint permissions. + +Candidate batches have shape ``(C, N, D_full)``, where the logical candidate +count ``C`` does not prescribe the host's physical batch size. ``valid_length`` +selects actual samples. Padding holds the final valid position and uses zero +time intervals; it is excluded from execution and recording. Extracting a row +copies its tensor payload and retains aligned identity, phase, factor, and +source-row metadata. + +An ``ExpertEpisode`` contains ``T`` commands actually sent and ``T+1`` +observations and timestamps, including the terminal observation. Episode +metadata accepts finite JSON values rather than live host objects. A validation +result accepts only a nonempty collection of passing checks; ``not_run``, +``failed``, and ``unavailable`` remain distinct rejection states. This structural +check does not generate physical evidence or certify the supplied checker. + +Persistence uses stable episode, candidate, scene, and commit IDs. The commit ID +remains fixed across write retries, while ``submission_id`` distinguishes each +submission and its delayed receipt. A persistence receipt is an input contract; +the package itself does not write or verify storage. + +.. autoclass:: SceneCase + :members: + +.. autoclass:: MotionSnapshot + :members: + +.. autoclass:: TrajectoryPhase + :members: + +.. autoclass:: TrajectoryTemplate + :members: + +.. autoclass:: CandidateIdentity + :members: + +.. autoclass:: CandidateTrajectoryBatch + :members: + +.. autoclass:: ValidationCheck + :members: + +.. autoclass:: ValidationResult + :members: + +.. autoclass:: ExpertEpisode + :members: + +.. autoclass:: CommitReceipt + :members: + +Configuration and Preflight +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``TrajectoryAugmentationCfg`` owns the local seed, provided-start declaration, +spatial and timing factors, and joint-geometry coverage limits. +``TrajectoryGenerationJobCfg`` adds the nested ``source``, ``planning``, +``execution``, ``reset``, ``validation``, ``collection``, and ``persistence`` +sections used by a standalone generation job. Both ``from_mapping`` methods +reject unknown fields and invalid types or ranges. ``validate_semantics`` +rechecks mutable configuration objects. + +The initial schema requires ``planning.batch_mode: env_rows``, +``execution.pool_mode: per_env_case``, ``execution.scheduler: full_batch``, +provided initial states, and synchronous persistence. It limits each candidate +to one rollout attempt. Write retries reuse the same episode and commit ID. +Other scheduling modes, overlapping planning and physics, and unsupported +enabled factors are rejected. Control periods remain owned by the host. + +Configuration IDs, including the default ``lerobot`` sink name, do not create +services. ``validate_capabilities`` requires explicit trusted source, validator, +profile, sink, and operator registries; configuration values are not imported +or evaluated. In particular, accepting a restoration profile ID or a supplied +``via_points`` capability does not implement physical restoration or an EEF +planner in this package. Deployment preflight must still verify those services. + +.. autoclass:: TrajectoryAugmentationCfg + :members: + :exclude-members: __init__, copy, replace, to_dict, validate + +.. autoclass:: TrajectoryGenerationJobCfg + :members: + :exclude-members: __init__, copy, replace, to_dict, validate + +Operators, Coverage, and Session +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``rotate_grasp_about_object_axis`` derives TCP pose candidates from a reference +grasp, rotating its position and orientation about an object-local axis through +the fixed object origin. The caller selects rotations appropriate to the +object's geometry, such as quarter turns for a cube, and replans each candidate. +The operator does not move the object or establish contact/IK validity. + +``joint_residual`` samples one smooth residual per explicitly permitted free +phase and preserves its endpoints and uncontrolled joints. ``retime`` rescales +permitted free phases and resamples on the supplied host control period, +remapping phase indices. Contact and hold durations are preserved. Generated +paths still need dynamic, collision, and task checks. +``validate_motion_limits`` supplies a sampled finite-difference speed and +acceleration check; it does not assert collision freedom or task success. + +``describe_trajectory`` computes phase-aligned normalized joint geometry and +elapsed-time descriptors from actual observations. ``CoverageIndex`` reserves +capacity before persistence and counts geometry only after confirmation. +Timing variants share geometry-family membership. This first descriptor uses +joint geometry; it does not provide EEF workspace coverage. + +``GenerationSession`` owns scene-case registration, stable candidate identity, +local random streams, bounded candidate and pending-episode payloads, collection +budgets, coverage reservations, and commit accounting. Its lifecycle accepts +planning results and measured episode evidence without stepping or resetting a +host. Duplicate receipts do not increment committed counts, and failed writes +release collection and coverage reservations before an explicit write retry. + +.. autofunction:: rotate_grasp_about_object_axis + +.. autofunction:: joint_residual + +.. autofunction:: retime + +.. autofunction:: validate_motion_limits + +.. autoclass:: TrajectoryDescriptor + :members: + +.. autofunction:: describe_trajectory + +.. autoclass:: CoverageIndex + :members: + +.. autoclass:: GenerationSession + :members: + +Implementation Modules +~~~~~~~~~~~~~~~~~~~~~~ + +The following module paths expose the same contracts and implementations. +The package import path above is convenient for callers combining them. + +.. currentmodule:: embodichain.lab.sim.motion.trajectory_augmentation.contracts + +.. autosummary:: + :nosignatures: + + SceneCase + MotionSnapshot + TrajectoryPhase + TrajectoryTemplate + CandidateIdentity + CandidateTrajectoryBatch + ValidationCheck + ValidationResult + ExpertEpisode + CommitReceipt + +.. currentmodule:: embodichain.lab.sim.motion.trajectory_augmentation.cfg + +.. autosummary:: + :nosignatures: + + TrajectoryAugmentationCfg + TrajectoryGenerationJobCfg + +.. currentmodule:: embodichain.lab.sim.motion.trajectory_augmentation.coverage + +.. autosummary:: + :nosignatures: + + TrajectoryDescriptor + describe_trajectory + CoverageIndex + +.. currentmodule:: embodichain.lab.sim.motion.trajectory_augmentation.operators + +.. autosummary:: + :nosignatures: + + rotate_grasp_about_object_axis + joint_residual + retime + validate_motion_limits + +.. currentmodule:: embodichain.lab.sim.motion.trajectory_augmentation.session + +.. autosummary:: + :nosignatures: + + GenerationSession diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.caches.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.caches.rst new file mode 100644 index 000000000..2e01b02a3 --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.caches.rst @@ -0,0 +1,7 @@ +embodichain.lab.sim.motion.workspace.caches +=========================================== + +.. automodule:: embodichain.lab.sim.motion.workspace.caches + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.configs.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.configs.rst new file mode 100644 index 000000000..5b75629c5 --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.configs.rst @@ -0,0 +1,7 @@ +embodichain.lab.sim.motion.workspace.configs +============================================ + +.. automodule:: embodichain.lab.sim.motion.workspace.configs + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.constraints.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.constraints.rst new file mode 100644 index 000000000..f834bd0ff --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.constraints.rst @@ -0,0 +1,7 @@ +embodichain.lab.sim.motion.workspace.constraints +================================================ + +.. automodule:: embodichain.lab.sim.motion.workspace.constraints + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.metrics.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.metrics.rst new file mode 100644 index 000000000..d98b840b3 --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.metrics.rst @@ -0,0 +1,7 @@ +embodichain.lab.sim.motion.workspace.metrics +============================================ + +.. automodule:: embodichain.lab.sim.motion.workspace.metrics + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.rst similarity index 60% rename from docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.rst rename to docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.rst index 3aee20004..3b56b93b1 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.rst @@ -1,7 +1,7 @@ -embodichain.lab.sim.workspace -============================= +embodichain.lab.sim.motion.workspace +==================================== -.. automodule:: embodichain.lab.sim.workspace +.. automodule:: embodichain.lab.sim.motion.workspace :members: :undoc-members: :show-inheritance: @@ -9,7 +9,7 @@ embodichain.lab.sim.workspace Runtime Sampling ---------------- -.. automodule:: embodichain.lab.sim.workspace.runtime +.. automodule:: embodichain.lab.sim.motion.workspace.runtime :members: :undoc-members: :show-inheritance: @@ -17,7 +17,7 @@ Runtime Sampling Runtime Configuration --------------------- -.. automodule:: embodichain.lab.sim.workspace.cfg +.. automodule:: embodichain.lab.sim.motion.workspace.cfg :members: :undoc-members: :show-inheritance: @@ -26,7 +26,7 @@ Runtime Configuration Offline Analysis ---------------- -.. automodule:: embodichain.lab.sim.workspace.analyzer +.. automodule:: embodichain.lab.sim.motion.workspace.analyzer :members: :undoc-members: :show-inheritance: @@ -34,7 +34,7 @@ Offline Analysis Workspace Components -------------------- -.. currentmodule:: embodichain.lab.sim.workspace +.. currentmodule:: embodichain.lab.sim.motion.workspace .. autosummary:: :toctree: . diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.samplers.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.samplers.rst new file mode 100644 index 000000000..08895aff4 --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.samplers.rst @@ -0,0 +1,7 @@ +embodichain.lab.sim.motion.workspace.samplers +============================================= + +.. automodule:: embodichain.lab.sim.motion.workspace.samplers + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.visualizers.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.visualizers.rst new file mode 100644 index 000000000..024fa874b --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.workspace.visualizers.rst @@ -0,0 +1,7 @@ +embodichain.lab.sim.motion.workspace.visualizers +================================================ + +.. automodule:: embodichain.lab.sim.motion.workspace.visualizers + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.rst index 716df56ef..a08de3cfb 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.rst @@ -9,10 +9,10 @@ Overview The ``sim`` package is EmbodiChain's simulation core. It is organized around the :class:`SimulationManager` (the DexSim scene handle), the scene-object hierarchy (lights, rigid/soft/cloth bodies, articulations, robots, gizmos, -constraints), the sensor suite (cameras, stereo cameras, contact sensors), IK -solvers and motion planners, the atomic-action motion-primitive layer, a -reusable workspace-analysis and sampling toolkit, and the shared configuration -types and utilities that wire all of these together. +constraints), the sensor suite (cameras, stereo cameras, contact sensors), the +``motion`` package for solvers, planners, workspace, and trajectory augmentation, +the atomic-action motion-primitive layer, and the shared configuration types and +utilities that wire all of these together. .. rubric:: Submodules @@ -28,10 +28,8 @@ types and utilities that wire all of these together. objects robots sensors - solvers - planners + motion atomic_actions - workspace types utility @@ -116,21 +114,13 @@ Robot Configurations embodichain.lab.sim.robots -Solvers -------- - -.. toctree:: - :maxdepth: 1 - - embodichain.lab.sim.solvers - -Planners --------- +Robot Motion +------------ .. toctree:: :maxdepth: 1 - embodichain.lab.sim.planners + embodichain.lab.sim.motion Atomic Actions -------------- @@ -140,14 +130,6 @@ Atomic Actions embodichain.lab.sim.atomic_actions -Robot Workspace ---------------- - -.. toctree:: - :maxdepth: 1 - - embodichain.lab.sim.workspace - Shared Types ------------ diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst index 1131a77f8..0f7e065b2 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst @@ -29,6 +29,11 @@ not write drive targets. Set this field to ``None`` to opt out or select prevents automatic recreation. ``GizmoCfg(ik_start_enabled=True)`` activates native IK on the first update with an open window, as used by the robot tutorial. +``simulation_time`` reports seconds successfully advanced through ``update`` +since manager construction. Collection hosts can use differences of this clock +for observation timestamps. Direct backend steps outside the manager, including +some object preparation setters, are not included. + .. rubric:: Classes .. autosummary:: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.caches.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.caches.rst deleted file mode 100644 index acd48527d..000000000 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.caches.rst +++ /dev/null @@ -1,7 +0,0 @@ -embodichain.lab.sim.workspace.caches -==================================== - -.. automodule:: embodichain.lab.sim.workspace.caches - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.configs.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.configs.rst deleted file mode 100644 index e3eb55ff2..000000000 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.configs.rst +++ /dev/null @@ -1,7 +0,0 @@ -embodichain.lab.sim.workspace.configs -===================================== - -.. automodule:: embodichain.lab.sim.workspace.configs - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.constraints.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.constraints.rst deleted file mode 100644 index 0734f3d9d..000000000 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.constraints.rst +++ /dev/null @@ -1,7 +0,0 @@ -embodichain.lab.sim.workspace.constraints -========================================= - -.. automodule:: embodichain.lab.sim.workspace.constraints - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.metrics.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.metrics.rst deleted file mode 100644 index 6142bbc60..000000000 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.metrics.rst +++ /dev/null @@ -1,7 +0,0 @@ -embodichain.lab.sim.workspace.metrics -===================================== - -.. automodule:: embodichain.lab.sim.workspace.metrics - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.samplers.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.samplers.rst deleted file mode 100644 index 80592a148..000000000 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.samplers.rst +++ /dev/null @@ -1,7 +0,0 @@ -embodichain.lab.sim.workspace.samplers -====================================== - -.. automodule:: embodichain.lab.sim.workspace.samplers - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.visualizers.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.visualizers.rst deleted file mode 100644 index a98937995..000000000 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.workspace.visualizers.rst +++ /dev/null @@ -1,7 +0,0 @@ -embodichain.lab.sim.workspace.visualizers -========================================= - -.. automodule:: embodichain.lab.sim.workspace.visualizers - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.trajectory_generation.rst b/docs/source/api_reference/embodichain/embodichain.lab.trajectory_generation.rst new file mode 100644 index 000000000..ebf53d168 --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.trajectory_generation.rst @@ -0,0 +1,390 @@ +embodichain.lab.trajectory_generation +===================================== + +.. automodule:: embodichain.lab.trajectory_generation + +The trajectory-generation integration layer prepares caller-provided fixed +scenes for repeated expert rollouts. It currently provides full-batch initial +state capture, restoration, verification, exclusive host ownership, qpos rollout +execution, a synchronous runner, and a LeRobot episode sink. Candidate operators +and generation accounting live in +:mod:`embodichain.lab.sim.motion.trajectory_augmentation`. +The planning adapter converts explicit EEF samples and validates supported free +qpos paths against the real environment batch. + +The synchronous runner supports free-motion references and an explicitly +bounded CPU-physics PickUp integration. The latter exports an offline atomic +compilation, checks full gripper and held-object geometry, records native +contacts at each physics substep, and persists accepted measured episodes. +AtomicActionRuntime recovery/feedback execution, contact-aware Gym collection, +a configuration registry, and the unified generation CLI remain separate work. +See :doc:`/overview/trajectory_generation` for the host lifecycle and integration +requirements. + +Initial-State Host +------------------ + +.. currentmodule:: embodichain.lab.trajectory_generation.initial_state + +.. autosummary:: + :nosignatures: + + InitialStateProfile + PreparedBatch + FixedSceneHost + +``InitialStateProfile`` supplies trusted Python callbacks for deterministic +task/controller preparation, fixed-condition identity, and task-specific initial +verification. A profile must cover controller history and task state beyond the +standard episode managers, and fixed physics, geometry, visual, sensor, and +control conditions beyond the physical snapshot. Profile IDs are descriptive; +they do not load callbacks from YAML or certify the implementation. + +``FixedSceneHost`` owns an entire simulator batch, with one ordered ``SceneCase`` +per physical row. ``acquire_case`` prepares and captures the initial states; +``restore_initial`` reconstructs them before a later rollout. ``PreparedBatch`` +identifies the current host epoch. Call ``assert_current`` before executing a +candidate; earlier or foreign bindings are rejected. Bindings do not transfer +old action plans, Task Program bridges, or other runtime state into a new epoch. +``snapshots(binding)`` copies planning inputs from the captured initial batch in +physical row order. ``initial_observation(binding)`` returns a copy of the Gym +first frame produced during preparation without querying observations again; +it returns ``None`` for a direct simulation host. + +When an ``EmbodiedEnv`` is supplied, the host uses its generation lease and +preparation API. Direct simulation and Gym ownership are mutually exclusive for +the same simulator batch. Gym resets the standard episode managers after +settling and before verification; initial observations and recording seeds are +published only after verification succeeds. The caller must freeze old episode +data and validation evidence before preparation discards its buffers. Closing the +host releases ownership; +it does not save or restore an episode. + +.. autoclass:: InitialStateProfile + :members: + :exclude-members: __init__, copy, replace, to_dict, validate + +.. autoclass:: PreparedBatch + :members: + +.. autoclass:: FixedSceneHost + :members: + +Simulation State Adapter +------------------------ + +.. currentmodule:: embodichain.lab.trajectory_generation.integrations.sim + +.. autosummary:: + :nosignatures: + + SimInitialState + SimInitialStateAdapter + +``SimInitialStateAdapter`` supports one configuration-owned fixed-base robot +and ordinary rigid objects spanning the full batch. The owned snapshot includes +local root/object poses, velocities, every movable joint in robot joint order, +position/velocity drive targets, and joint effort state. Tensor inputs are +detached and cloned; consumers must treat those owned tensors as read-only. + +The adapter checks topology, field names, shapes, finite values, SE(3) poses, +and joint limits before physical writes. Verification compares every supported +field with absolute tolerance ``atol`` and zero relative tolerance. It does not +snapshot fixed physical/visual properties, task/controller memory, contact +solver state, or pending external forces. Those boundaries require the trusted +profile and a supported episode initial state. +``tolerances_profile_id`` names this explicit tolerance policy for runner/job +matching; supplying the ID does not configure a tolerance or load a profile. + +Additional articulations, rigid-object groups, deformables, rigid constraints, +and robots whose articulation properties come from USD are rejected. Restoring +the root uses the existing articulation setter, which can advance the world by +1 ms; remaining joint and object writes follow that operation. A backend failure +can leave partial physical changes, so the owning host invalidates the binding +before restoration and only publishes a replacement after settling and +verification succeed. + +.. autoclass:: SimInitialState + :members: + +.. autoclass:: SimInitialStateAdapter + :members: + +Environment-Row Planning +------------------------ + +.. currentmodule:: embodichain.lab.trajectory_generation.integrations.planning + +.. autosummary:: + :nosignatures: + + EEFPath + EnvRowMotionPlanner + +``EEFPath`` owns explicit TCP poses ``(N,4,4)`` in the source row's local arena +frame, arrival intervals ``dt`` with an initial zero, candidate identity, source +row, and phase annotations. Optional ``solved_joint_targets`` have shape +``(N,D_control)`` in the selected control-part joint order. These samples preserve +a chosen IK branch; they are checked by FK without solving them again. +Float64 arrival intervals and supplied joint samples retain their precision. + +``EnvRowMotionPlanner`` requires one complete ``MotionSnapshot`` per real robot +row. It schedules candidates in rounds with at most one candidate per source row +and keeps backend calls at the real robot batch size. Input and output candidate +order, source-row identity, phase annotations, and arrival intervals stay aligned. +``plan_eef`` solves explicit EEF samples, propagates valid seeds, checks TCP FK +residuals, and returns a full-joint candidate batch plus per-candidate validation. +Failed rows hold their initial state and remain rejected by that paired result. + +``validate_qpos`` leaves the original path unchanged and densifies its segments +for collision checks, including phase boundaries. ``max_joint_step`` bounds the +sample spacing in joint coordinates and ``max_validation_samples`` caps samples +per candidate. Passing evidence is a sampled check, not continuous collision +detection, dynamic feasibility, or task success. +Collision-check scratch is allocated per real-batch round, bounded by ``B`` times +``max_validation_samples``, rather than retaining every candidate's dense path. + +Only fully annotated ``free`` motion with explicitly empty ``held_object_ids`` +is supported. Uncontrolled joints must stay at both the snapshot values and +robot configuration's initial values, matching the backend locked-joint model. +This requirement also applies when checking measured trajectories: a gripper +can drift under physics even if the candidate commands no gripper motion. +Measured locked-joint drift beyond the current ``1e-6`` comparison makes that +model unavailable for acceptance. +Contact/hold phases, unannotated samples, held-object sweeps, changing locked +joints, absent backend validation, +or excess sampling requirements return ``unavailable``. Invalid trajectories or +detected collisions return ``failed``. The caller maintains the same live robot +roots and collision world as the snapshots and supplies canonical dynamic +obstacle IDs; this adapter does not step or restore the scene. + +Measured qpos can be checked after execution while live root frames remain +fixed. Its first sample must match the snapshot within ``1e-6``. The obstacle +input represents one constant world for the whole path; the caller must +establish that it applies to the rollout. A single initial or final obstacle +pose does not reconstruct obstacle motion during execution. + +.. autoclass:: EEFPath + :members: + +.. autoclass:: EnvRowMotionPlanner + :members: + +Qpos Rollout Execution +---------------------- + +.. currentmodule:: embodichain.lab.trajectory_generation.execution + +.. autosummary:: + :nosignatures: + + QposRolloutExecutor + +``QposRolloutExecutor`` executes at most one candidate per physical row, with +``None`` for idle rows. A candidate's first sample is the initial state; +``valid_length - 1`` later samples are controller targets. The executor checks +case/epoch identity, full-joint order, initial state, active-joint permissions, +limits, the fixed clock, and estimated payload capacity before commands. +The payload limit includes tensors and metadata: execution reserves 64 KiB of +metadata capacity before commands and limits its own lineage, observation-key, +and validation metadata to 32 KiB, leaving room for later runner/storage evidence. +Oversized initial observations and changed schemas are rejected before the +executor creates its owned observation copy. + +Gym execution uses ``ControllerAction`` through the common demonstration loop. +The command observer copies actual full-joint position targets after controller +submission and before physics. The step observer consumes the returned Gym +observation without querying it again. Direct simulation uses the supplied +full-batch observation callback and an integer number of physics substeps per +command. Gym uses the authoritative ``env.step_dt``; both paths record differences +of ``SimulationManager.simulation_time`` and reject inconsistent clock advances. + +Returned ``ExpertEpisode`` values contain commands actually submitted and their +complete observation transitions. ``joint_positions`` always contains actual +full-joint measurements. Finished rows stop collecting frames while safe holds +allow other rows to finish; holds and padding are excluded from episode data. +Installing a hold clears velocity targets and joint effort without adding a +physics step. Execution remains serialized through task validation and evidence +freezing, including calls made from user-supplied validators. +The trusted task validator receives owned observations, actions, and timestamps +at the final batch boundary and must return its declared validation check. + +The executor records ``execution_complete`` and ``fixed_collision_world`` checks. +The latter compares robot roots and all rigid-object poses against the initial +snapshots at observation boundaries; it does not reconstruct motion between +those boundaries. Interrupted rollouts, target disagreement, changed observation +schemas, or runtime errors cannot produce accepted evidence. Rows without a +complete transition return ``None``. ``on_started`` fires after the first command +submission even if physics or observation then fails, so the rollout still +counts. Failure to install the final safe hold raises and stops collection. +``last_failures`` exposes an owned read-only mapping from candidate ID to the +most recent batch's runtime failure reason, including rows returning ``None``. +It resets on the next execution and retains at most one 512-character reason per +physical row, allowing the runner to preserve physics, observation, or cancellation +diagnostics without retaining rollout payloads. + +.. autoclass:: QposRolloutExecutor + :members: + +Synchronous Handwritten Runner +------------------------------ + +.. currentmodule:: embodichain.lab.trajectory_generation.runner + +.. autosummary:: + :nosignatures: + + MotionLimitsProfile + GenerationRunner + +``MotionLimitsProfile`` owns positive full-joint speed and acceleration vectors, +copied to CPU float64 in the robot's complete joint order. Its profile ID must +match the requested job policy. + +``GenerationRunner`` combines an explicit job configuration, ``FixedSceneHost``, +``EnvRowMotionPlanner``, qpos executor, synchronous LeRobot sink, and motion +limits. It checks integration IDs, robot ownership, joint layout, enabled +operators, and the host/sink clock before collection. ``run`` accepts one fixed +case and qpos template per physical row; the free-motion implementation +requires distinct case/initial-state pairs and the ``env_rows`` / ``full_batch`` +execution modes. + +With ``EnvRowMotionPlanner``, every host rigid-object physical UID must appear in both +``collision_world_entity_ids`` and ``dynamic_collision_entity_ids`` so the backend +receives each row's captured initial pose. Here dynamic means that backend poses +can be updated; a simulation body may still be static. A baked static world alone +cannot establish the prepared poses or different cases across rows. Different +relative layouts in a multi-row batch require per-environment backend worlds. +Semantic aliases are not connected in this runner. The trusted profile must +also certify collision representation for implicit +planes and other relevant geometry outside that rigid-object registry. +``executor.max_episode_bytes`` must fit both the sink limit and the job's pending +byte limit before rollout allocation. + +The runner proposes configured joint residuals and retiming, validates planned +paths and sampled motion limits, reserves episode capacity before rollout, and +restores the fixed initial batch between executed rounds. It rechecks measured +joint paths against the fixed collision snapshot and evaluates actual motion, +duration, task evidence, and motion limits before submitting an +accepted episode. Only confirmed sink receipts increase committed counts and +coverage. Failed writes can retry the same frozen episode within the configured +submission limit. + +Final evidence retains planning checks under ``planned_*`` IDs. Measured-path +collision results, ``actual_motion_limits``, and ``motion_quality`` are separate +acceptance checks; a passing planned path cannot replace their results. + +The runner is single-use and closes its supplied host and sink when ``run`` +exits. It returns a report with counters, coverage, audit, resolved settings, +package version, and ``target_reached``, and writes ``generation_report.json`` +under the sink root. Budget exhaustion can return without reaching the requested +target. Required integration, restoration, or persistence failures raise after +cleanup and report generation. + +``EnvRowMotionPlanner`` supports explicitly free, unloaded motion. The separate +``PickUpMotionValidator`` below adds changing gripper geometry, held-object +paths and physical contact gates for offline atomic qpos replay. The public +runner interface does not constitute full M1 runtime/source/host acceptance. + +.. autoclass:: MotionLimitsProfile + :members: + :exclude-members: __init__, copy, replace, to_dict, validate + +.. autoclass:: GenerationRunner + :members: + +Synchronous Episode Persistence +------------------------------- + +.. currentmodule:: embodichain.lab.trajectory_generation.sinks + +.. autosummary:: + :nosignatures: + + LeRobotEpisodeSink + +``LeRobotEpisodeSink`` writes each accepted ``ExpertEpisode`` into its own local +LeRobot dataset shard. Each shard contains exactly ``T`` causal training frames; +required sidecars retain the terminal observation, all ``T+1`` measured +timestamps, identity, validation, phases, and metadata. A collection manifest +maps committed episode IDs to their shards. + +``submit`` owns a CPU copy, writes and finalizes the dataset, reopens every frame +and required image, verifies sidecars, then replaces and reads back the manifest +before returning a confirmed ``CommitReceipt``. Invalid inputs raise before +episode writes; persistence failures return an unconfirmed receipt. Reusing a +confirmed commit ID with an identical payload only verifies existing artifacts; +it does not rewrite them or create another logical episode. +Retries can reuse readable sealed data, while incomplete uncommitted data is +rebuilt at the same shard path. Changed payloads under a known commit ID are +rejected. + +The sink accepts fixed-rate float32/float64 action vectors, supported numeric +observation vectors/matrices, and uint8 RGB images. Matrices flatten to C-order +vectors in LeRobot, retaining their source shape in ``observation_shapes`` and +original terminal matrix in the sidecar. Training timestamps use LeRobot's +relative fixed clock; the sidecar retains exact measured times and their origin. +Numeric readback checks the stored Parquet precision, while image readback uses +the LeRobot image decoder. +``max_episode_bytes`` limits raw tensor and metadata payload bytes for each +episode, not total process memory or collection disk size. Submission is +synchronous, so ``drain`` returns no deferred receipts and ``close`` releases the +writer lease after already completed writes. + +The sink requires a new or empty directory and one serialized caller. A confirmed +receipt means required artifacts were closed and successfully read back; process +restart and power-loss recovery are not implemented. It does not evaluate the +physical truth of supplied validation evidence or upload anything to the Hub. + +.. autoclass:: LeRobotEpisodeSink + :members: + +PickUp Source and Contact Validation +------------------------------------ + +.. currentmodule:: embodichain.lab.trajectory_generation.integrations.atomic + +.. autosummary:: + :nosignatures: + + export_pickup_templates + +The export accepts exactly a successful ``MoveEndEffector`` → ``PickUp`` +compilation. It preserves the atomic phase boundaries, expands passive mimic +geometry, makes the shared action boundary an explicit control-period hold, +and appends terminal hold commands. Only transit permits joint residuals. +It produces qpos replay templates; it does not execute AtomicActionRuntime or +commit the compilation's hypothetical symbolic effects. + +.. autofunction:: export_pickup_templates + +.. currentmodule:: embodichain.lab.trajectory_generation.integrations.contact + +.. autosummary:: + :nosignatures: + + PickUpContactProfile + PickUpMotionValidator + +Use the same ``PickUpMotionValidator`` as the Runner planner and the executor's +``contact_validator``. Supported geometry is an unscaled fixed-base URDF robot, +cuboid rigid objects, and the simulator's implicit ground. Planning uses +conservative hulls of URDF collision shapes and full-joint FK; actual validation +uses measured joints and measured object poses. The two-hop URDF self-collision +exclusions match the existing cuRobo structural policy. This remains sampled +validation, with additional native contact evidence at each CPU physics substep. + +The profile declares the target, support, finger links, fixed mounting links, +TCP, limited finger-contact entry region, penetration tolerance, and physical +hold thresholds. Unknown bodies, cross-row contact, buffer overflow, forbidden +contacts, missing finger contact, slipping and falling fail acceptance. The +profile does not enable contact or timing augmentation. Gym and arbitrary +mesh-shaped rigid objects are not supported by this integration. + +.. autoclass:: PickUpContactProfile + :members: + :show-inheritance: + +.. autoclass:: PickUpMotionValidator + :members: + :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.toolkits.rst b/docs/source/api_reference/embodichain/embodichain.toolkits.rst index e412f66a4..ba7513b41 100644 --- a/docs/source/api_reference/embodichain/embodichain.toolkits.rst +++ b/docs/source/api_reference/embodichain/embodichain.toolkits.rst @@ -1,8 +1,9 @@ embodichain.toolkits ==================== -The :mod:`embodichain.toolkits` package contains asset-preparation and -manipulation utilities that can be used independently of the simulation loop. +The :mod:`embodichain.toolkits` package contains standalone asset-preparation +and manipulation utilities. Robot motion capabilities are documented under +:mod:`embodichain.lab.sim.motion`. .. automodule:: embodichain.toolkits diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index 1c05b3565..af564f807 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -295,15 +295,6 @@ embodichain.lab.gym.envs.embodied_env EmbodiedEnvCfg EmbodiedEnv -embodichain.lab.task_program.integrations.simulation.handover ------------------------------------------------------------ - -.. currentmodule:: embodichain.lab.task_program.integrations.simulation.handover - -.. autosummary:: - - ConfiguredHandOverPoseProvider - embodichain.lab.gym.envs.managers.action_manager ------------------------------------------------ @@ -612,21 +603,6 @@ embodichain.lab.sim.atomic_actions.execution ExecutionStatus ExecutionTick -embodichain.lab.sim.atomic_actions.verification ------------------------------------------------- - -.. currentmodule:: embodichain.lab.sim.atomic_actions.verification - -.. autosummary:: - - EffectExpectationResult - EffectVerificationRequest - EffectVerificationResult - HeldObjectGuardRequest - HeldObjectGuardResult - PhaseEffectGateRequest - PhaseEffectGateResult - embodichain.lab.sim.atomic_actions.goals ---------------------------------------- @@ -829,99 +805,25 @@ embodichain.lab.sim.atomic_actions.transports EndpointCommandRouter EndpointCommandTransport -embodichain.lab.sim.objects.articulation ----------------------------------------- - -.. currentmodule:: embodichain.lab.sim.objects.articulation - -.. autosummary:: - - ArticulationData - Articulation - ArticulationJointKinematics - -embodichain.lab.sim.objects.cloth_object ----------------------------------------- - -.. currentmodule:: embodichain.lab.sim.objects.cloth_object - -.. autosummary:: - - ClothBodyData - ClothObject - ClothObjectCfg - -embodichain.lab.sim.objects.constraint --------------------------------------- - -.. currentmodule:: embodichain.lab.sim.objects.constraint - -.. autosummary:: - - RigidConstraint - -embodichain.lab.sim.objects.gizmo ---------------------------------- - -.. currentmodule:: embodichain.lab.sim.objects.gizmo - -Native robot targets use DexSim's controller with Newton IK by default. -Set ``GizmoCfg.ik_solver="embodichain"`` to reuse the robot control part's -configured solver, including PinkSolver; Viser uses the same solver adapter. - -.. autosummary:: - - Gizmo - GizmoCfg - create_robot_ik_gizmo_controller - -embodichain.lab.sim.objects.rigid_object ----------------------------------------- - -.. currentmodule:: embodichain.lab.sim.objects.rigid_object - -.. autosummary:: - - RigidBodyData - RigidObject - RigidObjectCfg - -embodichain.lab.sim.objects.rigid_object_group ----------------------------------------------- - -.. currentmodule:: embodichain.lab.sim.objects.rigid_object_group - -.. autosummary:: - - RigidBodyGroupData - RigidObjectGroup - RigidObjectGroupCfg - -embodichain.lab.sim.objects.robot ---------------------------------- - -.. currentmodule:: embodichain.lab.sim.objects.robot - -.. autosummary:: - - ControlGroup - Robot - -embodichain.lab.sim.objects.soft_object ---------------------------------------- +embodichain.lab.sim.atomic_actions.verification +------------------------------------------------ -.. currentmodule:: embodichain.lab.sim.objects.soft_object +.. currentmodule:: embodichain.lab.sim.atomic_actions.verification .. autosummary:: - SoftBodyData - SoftObject - SoftObjectCfg + EffectExpectationResult + EffectVerificationRequest + EffectVerificationResult + HeldObjectGuardRequest + HeldObjectGuardResult + PhaseEffectGateRequest + PhaseEffectGateResult -embodichain.lab.sim.planners.base_planner ------------------------------------------ +embodichain.lab.sim.motion.planners.base_planner +------------------------------------------------ -.. currentmodule:: embodichain.lab.sim.planners.base_planner +.. currentmodule:: embodichain.lab.sim.motion.planners.base_planner .. autosummary:: @@ -931,10 +833,10 @@ embodichain.lab.sim.planners.base_planner BasePlanner validate_plan_options -embodichain.lab.sim.planners.curobo.curobo_planner --------------------------------------------------- +embodichain.lab.sim.motion.planners.curobo.curobo_planner +--------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.planners.curobo.curobo_planner +.. currentmodule:: embodichain.lab.sim.motion.planners.curobo.curobo_planner .. autosummary:: @@ -944,20 +846,20 @@ embodichain.lab.sim.planners.curobo.curobo_planner CuroboPlannerCfg CuroboWorldCfg -embodichain.lab.sim.planners.curobo.curobo_yaml ------------------------------------------------ +embodichain.lab.sim.motion.planners.curobo.curobo_yaml +------------------------------------------------------ -.. currentmodule:: embodichain.lab.sim.planners.curobo.curobo_yaml +.. currentmodule:: embodichain.lab.sim.motion.planners.curobo.curobo_yaml .. autosummary:: generate_curobo_robot_yaml generate_curobo_world_yaml -embodichain.lab.sim.planners.motion_generator ---------------------------------------------- +embodichain.lab.sim.motion.planners.motion_generator +---------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.planners.motion_generator +.. currentmodule:: embodichain.lab.sim.motion.planners.motion_generator .. autosummary:: @@ -965,10 +867,10 @@ embodichain.lab.sim.planners.motion_generator MotionGenCfg MotionGenOptions -embodichain.lab.sim.planners.neural_planner -------------------------------------------- +embodichain.lab.sim.motion.planners.neural_planner +-------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.planners.neural_planner +.. currentmodule:: embodichain.lab.sim.motion.planners.neural_planner .. autosummary:: @@ -976,10 +878,10 @@ embodichain.lab.sim.planners.neural_planner NeuralPlannerCfg NeuralPlanOptions -embodichain.lab.sim.planners.toppra_planner -------------------------------------------- +embodichain.lab.sim.motion.planners.toppra_planner +-------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.planners.toppra_planner +.. currentmodule:: embodichain.lab.sim.motion.planners.toppra_planner .. autosummary:: @@ -987,10 +889,10 @@ embodichain.lab.sim.planners.toppra_planner ToppraPlannerCfg ToppraPlanOptions -embodichain.lab.sim.planners.utils ----------------------------------- +embodichain.lab.sim.motion.planners.utils +----------------------------------------- -.. currentmodule:: embodichain.lab.sim.planners.utils +.. currentmodule:: embodichain.lab.sim.motion.planners.utils .. autosummary:: @@ -1004,457 +906,555 @@ embodichain.lab.sim.planners.utils interpolate_xpos interpolate_xpos_batched -embodichain.lab.sim.robots.cobotmagic -------------------------------------- +embodichain.lab.sim.motion.solvers.neural_ik_solver +--------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.robots.cobotmagic +.. currentmodule:: embodichain.lab.sim.motion.solvers.neural_ik_solver .. autosummary:: - CobotMagicCfg + NeuralIKSolverCfg + NeuralIKSolver -embodichain.lab.sim.robots.dexforce_w1.hand_specs -------------------------------------------------- +embodichain.lab.sim.motion.solvers.null_space_posture_task +---------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.robots.dexforce_w1.hand_specs +.. currentmodule:: embodichain.lab.sim.motion.solvers.null_space_posture_task .. autosummary:: - W1HandSideSpec - W1HandSpec - get_default_w1_hand_version - get_w1_hand_spec - normalize_w1_hand_mappings + NullSpacePostureTask -embodichain.lab.sim.robots.dexforce_w1.specs --------------------------------------------- +embodichain.lab.sim.motion.solvers.pink_solver +---------------------------------------------- -.. currentmodule:: embodichain.lab.sim.robots.dexforce_w1.specs +.. currentmodule:: embodichain.lab.sim.motion.solvers.pink_solver .. autosummary:: - W1VersionSpec - get_w1_version_spec + PinkSolver + PinkSolverCfg -embodichain.lab.sim.robots.dexforce_w1.types --------------------------------------------- +embodichain.lab.sim.motion.solvers.srs_solver +--------------------------------------------- -.. currentmodule:: embodichain.lab.sim.robots.dexforce_w1.types +.. currentmodule:: embodichain.lab.sim.motion.solvers.srs_solver .. autosummary:: - DexforceW1Version - DexforceW1HandVersion - DexforceW1ArmSide - DexforceW1Type - DexforceW1HandBrand + SRSSolver + SRSSolverCfg -embodichain.lab.sim.robots.dexforce_w1.utils --------------------------------------------- +embodichain.lab.sim.motion.workspace.caches.cache_utils +------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.robots.dexforce_w1.utils +.. currentmodule:: embodichain.lab.sim.motion.workspace.caches.cache_utils .. autosummary:: - ChassisManager - TorsoManager - HeadManager - ArmManager - HandManager - EyesManager - build_dexforce_w1_assembly_urdf_cfg + clean_all_sessions + clean_session + format_size + get_cache_root + get_dir_size + list_sessions + main + show_session_info + show_total_size -embodichain.lab.sim.robots.dual_arm ------------------------------------ +embodichain.lab.sim.motion.workspace.caches.results_cache +--------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.robots.dual_arm +.. currentmodule:: embodichain.lab.sim.motion.workspace.caches.results_cache .. autosummary:: - DualArmRobotCfg - build_dual_arm_cfg - resolve_mounts + DEFAULT_RESULTS_CACHE_DIR + ResultsCache + compute_cache_key + serialize_results + deserialize_results -embodichain.lab.sim.robots.franka_panda ---------------------------------------- +embodichain.lab.sim.motion.workspace.constraints.base_constraint +---------------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.robots.franka_panda +.. currentmodule:: embodichain.lab.sim.motion.workspace.constraints.base_constraint .. autosummary:: - FrankaPandaCfg + IConstraintChecker + BaseConstraintChecker -embodichain.lab.sim.robots.ur_robot ------------------------------------ +embodichain.lab.sim.motion.workspace.constraints.workspace_constraint +--------------------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.robots.ur_robot +.. currentmodule:: embodichain.lab.sim.motion.workspace.constraints.workspace_constraint .. autosummary:: - URRobotCfg + WorkspaceConstraintChecker -embodichain.lab.sim.sensors.camera ----------------------------------- +embodichain.lab.sim.motion.workspace.samplers.base_sampler +---------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.sensors.camera +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.base_sampler .. autosummary:: - Camera - CameraCfg + ISampler + BaseSampler -embodichain.lab.sim.sim_manager -------------------------------- +embodichain.lab.sim.motion.workspace.samplers.gaussian_sampler +-------------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.sim_manager +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.gaussian_sampler .. autosummary:: - SIM_CACHE_DIR - MATERIAL_CACHE_DIR - CONVEX_DECOMP_DIR - REACHABLE_XPOS_DIR + GaussianSampler -embodichain.lab.task_program.semantics.calls ------------------------------------------------- +embodichain.lab.sim.motion.workspace.samplers.halton_sampler +------------------------------------------------------------ -.. currentmodule:: embodichain.lab.task_program.semantics.calls +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.halton_sampler .. autosummary:: - DeclarativeValue - HandOver - Pick - Place - PlaceRelationTarget - RegisteredSemanticCall - SemanticCallCatalog - SemanticCallDescriptor - SemanticCallSpec - SemanticPose - builtin_semantic_call_catalog - -.. automodule:: embodichain.lab.task_program.semantics.calls - :members: - :no-index: + HaltonSampler -embodichain.lab.task_program.semantics.integration ------------------------------------------------------- +embodichain.lab.sim.motion.workspace.samplers.importance_sampler +---------------------------------------------------------------- -.. currentmodule:: embodichain.lab.task_program.semantics.integration +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.importance_sampler .. autosummary:: - BoundSemanticCall - LinkedSemanticCall - PathPart - SceneEntityManifest - SceneManifest - SemanticDiagnostic - SemanticIntegrationManifest - SemanticValidationError + ImportanceSampler -.. automodule:: embodichain.lab.task_program.semantics.integration - :members: - :no-index: +embodichain.lab.sim.motion.workspace.samplers.iniform_sampler +------------------------------------------------------------- -embodichain.lab.task_program.semantics.effects --------------------------------------------------- +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.iniform_sampler -.. automodule:: embodichain.lab.task_program.semantics.effects - :members: - :no-index: +.. autosummary:: -embodichain.lab.task_program.semantics.evidence ---------------------------------------------------- + UniformSampler -.. automodule:: embodichain.lab.task_program.semantics.evidence - :members: - :no-index: +embodichain.lab.sim.motion.workspace.samplers.lhs_sampler +--------------------------------------------------------- -embodichain.lab.task_program.semantics.profiles ---------------------------------------------------- +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.lhs_sampler -.. currentmodule:: embodichain.lab.task_program.semantics.profiles +.. autosummary:: + + LatinHypercubeSampler + +embodichain.lab.sim.motion.workspace.samplers.random_sampler +------------------------------------------------------------ + +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.random_sampler .. autosummary:: - AmbiguousSkillBindingError - BoundRobotSkillProfile - ControlPartEndpoint - ControlPartEndpointAdapter - EffectAssurance - EndpointResolution - ProfileValidationError - ResourceEndpoint - ResourceEndpointAdapter - ResolvedRobotResource - ResolvedResourceEndpoint - ResolvedSkillBinding - ResourceBinding - ResourceClaim - RobotResource - RobotSkillProfile - SkillPolicyPreset - UnsupportedSkillError - WorkflowRecoveryPolicy + RandomSampler -embodichain.lab.task_program.semantics.scene ------------------------------------------------- +embodichain.lab.sim.motion.workspace.samplers.sobol_sampler +----------------------------------------------------------- -.. currentmodule:: embodichain.lab.task_program.semantics.scene +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.sobol_sampler .. autosummary:: - AmbiguousSceneAffordanceError - GRASP_AFFORDANCE_CAPABILITY - PLACE_IN_AFFORDANCE_CAPABILITY - PLACE_ON_AFFORDANCE_CAPABILITY - RegistrySceneProvider - SceneAffordanceRef - SceneArticulationRef - SceneCollisionRole - SceneCollisionWorldMode - SceneDynamics - SceneEntityRef - SceneEntityMetadata - SceneEntityRegistration - SceneEntityStateProvider - SceneGeometryProvider - SceneLinkRef - SceneObjectRef - SceneRegistry - UnsupportedSceneAffordanceError + SobolSampler -.. automodule:: embodichain.lab.task_program.semantics.scene - :members: - :no-index: +embodichain.lab.sim.motion.workspace.visualizers.axis_visualizer +---------------------------------------------------------------- -embodichain.lab.sim.solvers.neural_ik_solver --------------------------------------------- +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.axis_visualizer -.. currentmodule:: embodichain.lab.sim.solvers.neural_ik_solver +.. autosummary:: + + AxisVisualizer + +embodichain.lab.sim.motion.workspace.visualizers.base_visualizer +---------------------------------------------------------------- + +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.base_visualizer .. autosummary:: - NeuralIKSolverCfg - NeuralIKSolver + IVisualizer + BaseVisualizer -embodichain.lab.sim.solvers.null_space_posture_task ---------------------------------------------------- +embodichain.lab.sim.motion.workspace.visualizers.point_cloud_visualizer +----------------------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.solvers.null_space_posture_task +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.point_cloud_visualizer .. autosummary:: - NullSpacePostureTask + PointCloudVisualizer -embodichain.lab.sim.solvers.pink_solver ---------------------------------------- +embodichain.lab.sim.motion.workspace.visualizers.sphere_visualizer +------------------------------------------------------------------ -.. currentmodule:: embodichain.lab.sim.solvers.pink_solver +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.sphere_visualizer .. autosummary:: - PinkSolver - PinkSolverCfg + SphereVisualizer -embodichain.lab.sim.solvers.srs_solver --------------------------------------- +embodichain.lab.sim.motion.workspace.visualizers.visualizer_factory +------------------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.solvers.srs_solver +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.visualizer_factory .. autosummary:: - SRSSolver - SRSSolverCfg + VisualizerFactory + create_visualizer -embodichain.lab.sim.utility.render_utils +embodichain.lab.sim.motion.workspace.visualizers.voxel_visualizer +----------------------------------------------------------------- + +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.voxel_visualizer + +.. autosummary:: + + VoxelVisualizer + +embodichain.lab.sim.objects.articulation ---------------------------------------- -.. currentmodule:: embodichain.lab.sim.utility.render_utils +.. currentmodule:: embodichain.lab.sim.objects.articulation .. autosummary:: - select_default_renderer + ArticulationData + Articulation + ArticulationJointKinematics -embodichain.lab.sim.workspace.caches.cache_utils ------------------------------------------------- +embodichain.lab.sim.objects.cloth_object +---------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.caches.cache_utils +.. currentmodule:: embodichain.lab.sim.objects.cloth_object .. autosummary:: - clean_all_sessions - clean_session - format_size - get_cache_root - get_dir_size - list_sessions - main - show_session_info - show_total_size + ClothBodyData + ClothObject + ClothObjectCfg -embodichain.lab.sim.workspace.caches.results_cache --------------------------------------------------- +embodichain.lab.sim.objects.constraint +-------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.caches.results_cache +.. currentmodule:: embodichain.lab.sim.objects.constraint .. autosummary:: - DEFAULT_RESULTS_CACHE_DIR - ResultsCache - compute_cache_key - serialize_results - deserialize_results + RigidConstraint -embodichain.lab.sim.workspace.constraints.base_constraint ---------------------------------------------------------- +embodichain.lab.sim.objects.gizmo +--------------------------------- + +.. currentmodule:: embodichain.lab.sim.objects.gizmo -.. currentmodule:: embodichain.lab.sim.workspace.constraints.base_constraint +Native robot targets use DexSim's controller with Newton IK by default. +Set ``GizmoCfg.ik_solver="embodichain"`` to reuse the robot control part's +configured solver, including PinkSolver; Viser uses the same solver adapter. .. autosummary:: - IConstraintChecker - BaseConstraintChecker + Gizmo + GizmoCfg + create_robot_ik_gizmo_controller -embodichain.lab.sim.workspace.constraints.workspace_constraint --------------------------------------------------------------- +embodichain.lab.sim.objects.rigid_object +---------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.constraints.workspace_constraint +.. currentmodule:: embodichain.lab.sim.objects.rigid_object .. autosummary:: - WorkspaceConstraintChecker + RigidBodyData + RigidObject + RigidObjectCfg -embodichain.lab.sim.workspace.samplers.base_sampler ---------------------------------------------------- +embodichain.lab.sim.objects.rigid_object_group +---------------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.samplers.base_sampler +.. currentmodule:: embodichain.lab.sim.objects.rigid_object_group .. autosummary:: - ISampler - BaseSampler + RigidBodyGroupData + RigidObjectGroup + RigidObjectGroupCfg -embodichain.lab.sim.workspace.samplers.gaussian_sampler -------------------------------------------------------- +embodichain.lab.sim.objects.robot +--------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.samplers.gaussian_sampler +.. currentmodule:: embodichain.lab.sim.objects.robot .. autosummary:: - GaussianSampler + ControlGroup + Robot -embodichain.lab.sim.workspace.samplers.halton_sampler ------------------------------------------------------ +embodichain.lab.sim.objects.soft_object +--------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.samplers.halton_sampler +.. currentmodule:: embodichain.lab.sim.objects.soft_object .. autosummary:: - HaltonSampler + SoftBodyData + SoftObject + SoftObjectCfg -embodichain.lab.sim.workspace.samplers.importance_sampler ---------------------------------------------------------- +embodichain.lab.sim.robots.cobotmagic +------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.samplers.importance_sampler +.. currentmodule:: embodichain.lab.sim.robots.cobotmagic .. autosummary:: - ImportanceSampler + CobotMagicCfg -embodichain.lab.sim.workspace.samplers.iniform_sampler ------------------------------------------------------- +embodichain.lab.sim.robots.dexforce_w1.hand_specs +------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.samplers.iniform_sampler +.. currentmodule:: embodichain.lab.sim.robots.dexforce_w1.hand_specs .. autosummary:: - UniformSampler + W1HandSideSpec + W1HandSpec + get_default_w1_hand_version + get_w1_hand_spec + normalize_w1_hand_mappings -embodichain.lab.sim.workspace.samplers.lhs_sampler --------------------------------------------------- +embodichain.lab.sim.robots.dexforce_w1.specs +-------------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.samplers.lhs_sampler +.. currentmodule:: embodichain.lab.sim.robots.dexforce_w1.specs .. autosummary:: - LatinHypercubeSampler + W1VersionSpec + get_w1_version_spec -embodichain.lab.sim.workspace.samplers.random_sampler ------------------------------------------------------ +embodichain.lab.sim.robots.dexforce_w1.types +-------------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.samplers.random_sampler +.. currentmodule:: embodichain.lab.sim.robots.dexforce_w1.types .. autosummary:: - RandomSampler + DexforceW1Version + DexforceW1HandVersion + DexforceW1ArmSide + DexforceW1Type + DexforceW1HandBrand -embodichain.lab.sim.workspace.samplers.sobol_sampler ----------------------------------------------------- +embodichain.lab.sim.robots.dexforce_w1.utils +-------------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.samplers.sobol_sampler +.. currentmodule:: embodichain.lab.sim.robots.dexforce_w1.utils .. autosummary:: - SobolSampler + ChassisManager + TorsoManager + HeadManager + ArmManager + HandManager + EyesManager + build_dexforce_w1_assembly_urdf_cfg -embodichain.lab.sim.workspace.visualizers.axis_visualizer ---------------------------------------------------------- +embodichain.lab.sim.robots.dual_arm +----------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.visualizers.axis_visualizer +.. currentmodule:: embodichain.lab.sim.robots.dual_arm .. autosummary:: - AxisVisualizer + DualArmRobotCfg + build_dual_arm_cfg + resolve_mounts -embodichain.lab.sim.workspace.visualizers.base_visualizer ---------------------------------------------------------- +embodichain.lab.sim.robots.franka_panda +--------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.visualizers.base_visualizer +.. currentmodule:: embodichain.lab.sim.robots.franka_panda .. autosummary:: - IVisualizer - BaseVisualizer + FrankaPandaCfg -embodichain.lab.sim.workspace.visualizers.point_cloud_visualizer ----------------------------------------------------------------- +embodichain.lab.sim.robots.ur_robot +----------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.visualizers.point_cloud_visualizer +.. currentmodule:: embodichain.lab.sim.robots.ur_robot .. autosummary:: - PointCloudVisualizer + URRobotCfg + +embodichain.lab.sim.sensors.camera +---------------------------------- + +.. currentmodule:: embodichain.lab.sim.sensors.camera + +.. autosummary:: + + Camera + CameraCfg + +embodichain.lab.sim.sim_manager +------------------------------- + +.. currentmodule:: embodichain.lab.sim.sim_manager + +.. autosummary:: + + SIM_CACHE_DIR + MATERIAL_CACHE_DIR + CONVEX_DECOMP_DIR + REACHABLE_XPOS_DIR + +embodichain.lab.sim.utility.render_utils +---------------------------------------- -embodichain.lab.sim.workspace.visualizers.sphere_visualizer +.. currentmodule:: embodichain.lab.sim.utility.render_utils + +.. autosummary:: + + select_default_renderer + +embodichain.lab.task_program.integrations.simulation.handover ----------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.workspace.visualizers.sphere_visualizer +.. currentmodule:: embodichain.lab.task_program.integrations.simulation.handover .. autosummary:: - SphereVisualizer + ConfiguredHandOverPoseProvider -embodichain.lab.sim.workspace.visualizers.visualizer_factory ------------------------------------------------------------- +embodichain.lab.task_program.semantics.calls +------------------------------------------------ -.. currentmodule:: embodichain.lab.sim.workspace.visualizers.visualizer_factory +.. currentmodule:: embodichain.lab.task_program.semantics.calls .. autosummary:: - VisualizerFactory - create_visualizer + DeclarativeValue + HandOver + Pick + Place + PlaceRelationTarget + RegisteredSemanticCall + SemanticCallCatalog + SemanticCallDescriptor + SemanticCallSpec + SemanticPose + builtin_semantic_call_catalog -embodichain.lab.sim.workspace.visualizers.voxel_visualizer ----------------------------------------------------------- +.. automodule:: embodichain.lab.task_program.semantics.calls + :members: + :no-index: + +embodichain.lab.task_program.semantics.effects +-------------------------------------------------- + +.. automodule:: embodichain.lab.task_program.semantics.effects + :members: + :no-index: + +embodichain.lab.task_program.semantics.evidence +--------------------------------------------------- + +.. automodule:: embodichain.lab.task_program.semantics.evidence + :members: + :no-index: + +embodichain.lab.task_program.semantics.integration +------------------------------------------------------ -.. currentmodule:: embodichain.lab.sim.workspace.visualizers.voxel_visualizer +.. currentmodule:: embodichain.lab.task_program.semantics.integration .. autosummary:: - VoxelVisualizer + BoundSemanticCall + LinkedSemanticCall + PathPart + SceneEntityManifest + SceneManifest + SemanticDiagnostic + SemanticIntegrationManifest + SemanticValidationError + +.. automodule:: embodichain.lab.task_program.semantics.integration + :members: + :no-index: + +embodichain.lab.task_program.semantics.profiles +--------------------------------------------------- + +.. currentmodule:: embodichain.lab.task_program.semantics.profiles + +.. autosummary:: + + AmbiguousSkillBindingError + BoundRobotSkillProfile + ControlPartEndpoint + ControlPartEndpointAdapter + EffectAssurance + EndpointResolution + ProfileValidationError + ResourceEndpoint + ResourceEndpointAdapter + ResolvedRobotResource + ResolvedResourceEndpoint + ResolvedSkillBinding + ResourceBinding + ResourceClaim + RobotResource + RobotSkillProfile + SkillPolicyPreset + UnsupportedSkillError + WorkflowRecoveryPolicy + +embodichain.lab.task_program.semantics.scene +------------------------------------------------ + +.. currentmodule:: embodichain.lab.task_program.semantics.scene + +.. autosummary:: + + AmbiguousSceneAffordanceError + GRASP_AFFORDANCE_CAPABILITY + PLACE_IN_AFFORDANCE_CAPABILITY + PLACE_ON_AFFORDANCE_CAPABILITY + RegistrySceneProvider + SceneAffordanceRef + SceneArticulationRef + SceneCollisionRole + SceneCollisionWorldMode + SceneDynamics + SceneEntityRef + SceneEntityMetadata + SceneEntityRegistration + SceneEntityStateProvider + SceneGeometryProvider + SceneLinkRef + SceneObjectRef + SceneRegistry + UnsupportedSceneAffordanceError + +.. automodule:: embodichain.lab.task_program.semantics.scene + :members: + :no-index: embodichain.lab.visualization.backends -------------------------------------- diff --git a/docs/source/features/workspace_analyzer/caches.md b/docs/source/features/workspace_analyzer/caches.md index 5523aaa62..5f5b6e9ff 100644 --- a/docs/source/features/workspace_analyzer/caches.md +++ b/docs/source/features/workspace_analyzer/caches.md @@ -31,7 +31,7 @@ The caches module provides simple caching mechanisms for workspace analysis data Fast in-memory caching for medium-sized datasets: ```python -from embodichain.lab.sim.workspace.caches import MemoryCache +from embodichain.lab.sim.motion.workspace.caches import MemoryCache # Create memory cache cache = MemoryCache() @@ -49,7 +49,7 @@ all_data = cache.get_all() Persistent disk caching for large datasets: ```python -from embodichain.lab.sim.workspace.caches import DiskCache +from embodichain.lab.sim.motion.workspace.caches import DiskCache # Create disk cache cache = DiskCache(save_dir="./my_cache") @@ -68,7 +68,7 @@ all_data = cache.get_all() Simplified factory pattern for creating caches: ```python -from embodichain.lab.sim.workspace.caches import CacheManager +from embodichain.lab.sim.motion.workspace.caches import CacheManager # Create memory cache memory_cache = CacheManager.create_cache("memory") @@ -83,7 +83,7 @@ disk_cache = CacheManager.create_cache("disk", save_dir="./cache") ```python import numpy as np -from embodichain.lab.sim.workspace.caches import MemoryCache, DiskCache +from embodichain.lab.sim.motion.workspace.caches import MemoryCache, DiskCache # Memory cache example memory_cache = MemoryCache() @@ -196,7 +196,7 @@ cache = CacheManager.create_cache("memory") # or "disk" ```python # Simplest usage -from embodichain.lab.sim.workspace.caches import MemoryCache, DiskCache +from embodichain.lab.sim.motion.workspace.caches import MemoryCache, DiskCache # For small datasets cache = MemoryCache() diff --git a/docs/source/features/workspace_analyzer/cli.md b/docs/source/features/workspace_analyzer/cli.md index 2e53a2de6..cbb977772 100644 --- a/docs/source/features/workspace_analyzer/cli.md +++ b/docs/source/features/workspace_analyzer/cli.md @@ -194,7 +194,7 @@ configurations are converted to poses using each environment's current robot base: ```python -from embodichain.lab.sim.workspace import RobotWorkspace +from embodichain.lab.sim.motion.workspace import RobotWorkspace workspace = RobotWorkspace.from_cache(entry, device="cuda") indices = workspace.sample_indices(16, strategy="voxel_uniform") @@ -207,7 +207,7 @@ integration. To look up an entry by its inputs from Python, use the analyzer's cache key: ```python -from embodichain.lab.sim.workspace.caches import ( +from embodichain.lab.sim.motion.workspace.caches import ( ResultsCache, compute_cache_key, ) # metadata = analyzer._build_cache_key_metadata(num_samples) # same inputs diff --git a/docs/source/features/workspace_analyzer/configs.md b/docs/source/features/workspace_analyzer/configs.md index 90b2f2cc6..2bc440bca 100644 --- a/docs/source/features/workspace_analyzer/configs.md +++ b/docs/source/features/workspace_analyzer/configs.md @@ -37,7 +37,7 @@ The configs module enables fine-grained control over: Controls how the robot's joint space is sampled for workspace analysis. ```python -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( SamplingConfig, SamplingStrategy ) @@ -103,7 +103,7 @@ importance_config = SamplingConfig( Manages caching of analysis results for improved performance. ```python -from embodichain.lab.sim.workspace.configs import CacheConfig +from embodichain.lab.sim.motion.workspace.configs import CacheConfig from pathlib import Path # Basic caching configuration @@ -139,7 +139,7 @@ performance_cache = CacheConfig( Defines which workspace metrics to compute and their parameters. ```python -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( MetricConfig, MetricType, ReachabilityConfig, ManipulabilityConfig, DensityConfig ) @@ -216,7 +216,7 @@ density_config = DensityConfig( Controls the appearance and behavior of workspace visualizations. ```python -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( VisualizationConfig, VisualizationType ) @@ -256,7 +256,7 @@ sphere_viz_config = VisualizationConfig( Defines workspace bounds, exclusion zones, and physical constraints. ```python -from embodichain.lab.sim.workspace.configs import DimensionConstraint +from embodichain.lab.sim.motion.workspace.configs import DimensionConstraint import numpy as np # Basic workspace bounds @@ -298,7 +298,7 @@ complex_constraints = DimensionConstraint( ### Basic Workspace Analysis Setup ```python -from embodichain.lab.sim.workspace.configs import * +from embodichain.lab.sim.motion.workspace.configs import * # Complete configuration for basic analysis sampling_config = SamplingConfig( @@ -597,7 +597,7 @@ Configuration for workspace bounds and constraints. ```python # Simplest possible configuration -from embodichain.lab.sim.workspace.configs import * +from embodichain.lab.sim.motion.workspace.configs import * config = SamplingConfig() # Uses defaults # Result: uniform sampling, 1000 samples, grid resolution 10 diff --git a/docs/source/features/workspace_analyzer/constraints.md b/docs/source/features/workspace_analyzer/constraints.md index c2af9c005..273dd1104 100644 --- a/docs/source/features/workspace_analyzer/constraints.md +++ b/docs/source/features/workspace_analyzer/constraints.md @@ -45,7 +45,7 @@ Define workspace boundaries using min/max bounds: ```python import numpy as np -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Basic workspace boundaries checker = WorkspaceConstraintChecker( @@ -61,7 +61,7 @@ Avoid specified rectangular obstacle regions (simple collision avoidance): ```python import numpy as np -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Workspace with exclusion zones checker = WorkspaceConstraintChecker( @@ -82,8 +82,8 @@ Use DimensionConstraint for easy configuration: ```python import numpy as np -from embodichain.lab.sim.workspace.configs.dimension_constraint import DimensionConstraint -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.configs.dimension_constraint import DimensionConstraint +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Create configuration config = DimensionConstraint( @@ -105,7 +105,7 @@ checker = WorkspaceConstraintChecker.from_config(config) ```python import numpy as np -from embodichain.lab.sim.workspace.configs.dimension_constraint import DimensionConstraint +from embodichain.lab.sim.motion.workspace.configs.dimension_constraint import DimensionConstraint # TODO: These parameters exist in DimensionConstraint but are not implemented config = DimensionConstraint( @@ -135,7 +135,7 @@ config = DimensionConstraint( ```python import numpy as np -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Create constraint checker checker = WorkspaceConstraintChecker( @@ -165,7 +165,7 @@ print(f"Filtered points shape: {valid_points.shape}") # (2, 3) ```python import numpy as np -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Workspace with exclusion zones checker = WorkspaceConstraintChecker( @@ -197,8 +197,8 @@ print(f"Safe points: {len(safe_points)} out of {len(test_points)}") ```python import numpy as np -from embodichain.lab.sim.workspace.configs.dimension_constraint import DimensionConstraint -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.configs.dimension_constraint import DimensionConstraint +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Create configuration object config = DimensionConstraint( @@ -227,7 +227,7 @@ print(f"Total exclusion zones: {checker.get_num_exclude_zones()}") ```python import numpy as np -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Add and remove obstacles dynamically checker = WorkspaceConstraintChecker( @@ -252,8 +252,8 @@ print(f"After clearing: {checker.get_num_exclude_zones()} zones") ```python import numpy as np -from embodichain.lab.sim.workspace.configs.dimension_constraint import DimensionConstraint -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.configs.dimension_constraint import DimensionConstraint +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Define workspace constraints constraints = DimensionConstraint( @@ -282,7 +282,7 @@ print(f"Coverage: {len(valid_points)/len(workspace_points)*100:.1f}%") ```python import numpy as np -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Conservative workspace bounds def create_safe_bounds(robot_reach, safety_margin=0.1): @@ -368,7 +368,7 @@ Configuration dataclass for constraints. ```python import numpy as np -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Simple boundary constraints checker = WorkspaceConstraintChecker( @@ -381,7 +381,7 @@ checker = WorkspaceConstraintChecker( ```python import numpy as np -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Workspace with exclusion zones checker = WorkspaceConstraintChecker( @@ -398,8 +398,8 @@ checker = WorkspaceConstraintChecker( ```python import numpy as np -from embodichain.lab.sim.workspace.configs.dimension_constraint import DimensionConstraint -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.configs.dimension_constraint import DimensionConstraint +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Create config config = DimensionConstraint( diff --git a/docs/source/features/workspace_analyzer/index.rst b/docs/source/features/workspace_analyzer/index.rst index 2ce6b2130..508b18d42 100644 --- a/docs/source/features/workspace_analyzer/index.rst +++ b/docs/source/features/workspace_analyzer/index.rst @@ -1,7 +1,7 @@ Robot Workspace =============== -The ``embodichain.lab.sim.workspace`` package provides offline workspace +The ``embodichain.lab.sim.motion.workspace`` package provides offline workspace analysis, persistent cache management, visualization, and runtime sampling for task-environment randomization. diff --git a/docs/source/features/workspace_analyzer/metrics.md b/docs/source/features/workspace_analyzer/metrics.md index 0f63fed84..400800671 100644 --- a/docs/source/features/workspace_analyzer/metrics.md +++ b/docs/source/features/workspace_analyzer/metrics.md @@ -27,7 +27,7 @@ Computes workspace volume and coverage using voxel-based analysis. ```python import numpy as np -from embodichain.lab.sim.workspace.metrics import ReachabilityMetric +from embodichain.lab.sim.motion.workspace.metrics import ReachabilityMetric # Basic usage workspace_points = np.random.uniform(-1, 1, size=(1000, 3)) @@ -46,7 +46,7 @@ Analyzes dexterity using distance-based heuristic or Yoshikawa index (with Jacob ```python import numpy as np -from embodichain.lab.sim.workspace.metrics import ManipulabilityMetric +from embodichain.lab.sim.motion.workspace.metrics import ManipulabilityMetric # Basic usage (heuristic method) workspace_points = np.random.uniform(-1, 1, size=(1000, 3)) @@ -64,7 +64,7 @@ Computes local point density using radius-based neighborhood analysis. ```python import numpy as np -from embodichain.lab.sim.workspace.metrics import DensityMetric +from embodichain.lab.sim.motion.workspace.metrics import DensityMetric # Basic usage workspace_points = np.random.uniform(-1, 1, size=(1000, 3)) @@ -82,7 +82,7 @@ print(f"Mean density: {results['mean_density']:.2f}") ```python import numpy as np -from embodichain.lab.sim.workspace.metrics import ( +from embodichain.lab.sim.motion.workspace.metrics import ( ReachabilityMetric, ManipulabilityMetric, DensityMetric ) @@ -102,7 +102,7 @@ print(f"Mean density: {density_results['mean_density']:.2f}") ### With Custom Configuration ```python -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( ReachabilityConfig, DensityConfig ) @@ -122,7 +122,7 @@ density_results = density.compute(workspace_points) All metrics can be customized using configuration classes: ```python -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( ReachabilityConfig, ManipulabilityConfig, DensityConfig ) @@ -163,7 +163,7 @@ manipulability = ManipulabilityMetric() density = DensityMetric() # Custom configuration -from embodichain.lab.sim.workspace.configs import ReachabilityConfig +from embodichain.lab.sim.motion.workspace.configs import ReachabilityConfig reach_config = ReachabilityConfig(voxel_size=0.02, compute_coverage=True) reachability = ReachabilityMetric(reach_config) ``` diff --git a/docs/source/features/workspace_analyzer/runtime.md b/docs/source/features/workspace_analyzer/runtime.md index 8fbca8335..35340ccac 100644 --- a/docs/source/features/workspace_analyzer/runtime.md +++ b/docs/source/features/workspace_analyzer/runtime.md @@ -1,6 +1,6 @@ # Runtime Workspace Sampling -The `embodichain.lab.sim.workspace` package can reuse a cached workspace during +The `embodichain.lab.sim.motion.workspace` package can reuse a cached workspace during environment resets. Runtime sampling selects a reachable joint configuration from the cache and asks the target `Robot` to recompute forward kinematics for each environment. This accounts for the current robot-base pose instead of @@ -14,7 +14,7 @@ that a collision-free trajectory exists from the robot's current state. Workspace caches are configured per robot control part: ```python -from embodichain.lab.sim.workspace import RobotWorkspaceCfg +from embodichain.lab.sim.motion.workspace import RobotWorkspaceCfg workspace_cfg = { "left_arm": RobotWorkspaceCfg( @@ -103,7 +103,7 @@ their previous object pose. ## Package layout ```text -embodichain/lab/sim/workspace/ +embodichain/lab/sim/motion/workspace/ ├── runtime.py # RobotWorkspace and WorkspaceSample ├── cfg.py # RobotWorkspaceCfg ├── analyzer.py # WorkspaceAnalyzer diff --git a/docs/source/features/workspace_analyzer/samplers.md b/docs/source/features/workspace_analyzer/samplers.md index 70087aa23..7285f3523 100644 --- a/docs/source/features/workspace_analyzer/samplers.md +++ b/docs/source/features/workspace_analyzer/samplers.md @@ -10,7 +10,7 @@ grid may return a nearby grid-sized count). ```python import torch -from embodichain.lab.sim.workspace.samplers import UniformSampler +from embodichain.lab.sim.motion.workspace.samplers import UniformSampler bounds = torch.tensor( [ @@ -49,8 +49,8 @@ sampling; it does not apply a spherical geometric constraint. Use `create_sampler` when the strategy comes from configuration: ```python -from embodichain.lab.sim.workspace.configs import SamplingStrategy -from embodichain.lab.sim.workspace.samplers import create_sampler +from embodichain.lab.sim.motion.workspace.configs import SamplingStrategy +from embodichain.lab.sim.motion.workspace.samplers import create_sampler sampler = create_sampler( SamplingStrategy.SOBOL, @@ -89,7 +89,7 @@ Select a strategy through `SamplingConfig`; `WorkspaceAnalyzer` creates the matching sampler and uses it for joint- and Cartesian-space sampling: ```python -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( SamplingConfig, SamplingStrategy, ) diff --git a/docs/source/features/workspace_analyzer/visualizers.md b/docs/source/features/workspace_analyzer/visualizers.md index 1f69ff4e0..8733cb70e 100644 --- a/docs/source/features/workspace_analyzer/visualizers.md +++ b/docs/source/features/workspace_analyzer/visualizers.md @@ -26,7 +26,7 @@ The visualizers module enables: Fast rendering for large datasets. ```python -from embodichain.lab.sim.workspace.visualizers import PointCloudVisualizer +from embodichain.lab.sim.motion.workspace.visualizers import PointCloudVisualizer visualizer = PointCloudVisualizer(backend='open3d', point_size=2.0) ``` @@ -38,7 +38,7 @@ visualizer = PointCloudVisualizer(backend='open3d', point_size=2.0) Volumetric representation for occupancy maps. ```python -from embodichain.lab.sim.workspace.visualizers import VoxelVisualizer +from embodichain.lab.sim.motion.workspace.visualizers import VoxelVisualizer visualizer = VoxelVisualizer(backend='open3d', voxel_size=0.01) ``` @@ -50,7 +50,7 @@ visualizer = VoxelVisualizer(backend='open3d', voxel_size=0.01) Smooth visualization for reachability zones. ```python -from embodichain.lab.sim.workspace.visualizers import SphereVisualizer +from embodichain.lab.sim.motion.workspace.visualizers import SphereVisualizer visualizer = SphereVisualizer(backend='open3d', sphere_radius=0.005) ``` @@ -62,7 +62,7 @@ visualizer = SphereVisualizer(backend='open3d', sphere_radius=0.005) Coordinate frame visualization for poses and transformations. ```python -from embodichain.lab.sim.workspace.visualizers import AxisVisualizer +from embodichain.lab.sim.motion.workspace.visualizers import AxisVisualizer visualizer = AxisVisualizer(backend='sim_manager', sim_manager=sim, axis_length=0.15) ``` @@ -75,7 +75,7 @@ visualizer = AxisVisualizer(backend='sim_manager', sim_manager=sim, axis_length= ```python import numpy as np -from embodichain.lab.sim.workspace.visualizers import create_visualizer +from embodichain.lab.sim.motion.workspace.visualizers import create_visualizer # Generate workspace data points = np.random.rand(1000, 3) * 2 - 1 # Random points in [-1, 1]³ @@ -103,7 +103,7 @@ visualizer.visualize(points, colors) ### Factory Pattern ```python -from embodichain.lab.sim.workspace.visualizers import VisualizerFactory +from embodichain.lab.sim.motion.workspace.visualizers import VisualizerFactory factory = VisualizerFactory() visualizer = factory.create_visualizer('voxel', backend='open3d', voxel_size=0.02) @@ -140,7 +140,7 @@ axis_viz = factory.create_visualizer('axis', backend='sim_manager', sim_manager= **Quick Creation**: ```python -from embodichain.lab.sim.workspace.visualizers import create_visualizer +from embodichain.lab.sim.motion.workspace.visualizers import create_visualizer import numpy as np # Create any visualizer diff --git a/docs/source/features/workspace_analyzer/workspace_analyzer.md b/docs/source/features/workspace_analyzer/workspace_analyzer.md index 2e49ff0e0..025f449f0 100644 --- a/docs/source/features/workspace_analyzer/workspace_analyzer.md +++ b/docs/source/features/workspace_analyzer/workspace_analyzer.md @@ -3,7 +3,7 @@ The Workspace Analyzer is a comprehensive tool in EmbodiChain for analyzing robot workspace characteristics, reachability, and performance metrics. It provides multiple analysis modes with advanced sampling strategies, caching mechanisms, and visualization capabilities. The analyzer and runtime sampling APIs are unified under -`embodichain.lab.sim.workspace`. See [Runtime Workspace Sampling](runtime.md) +`embodichain.lab.sim.motion.workspace`. See [Runtime Workspace Sampling](runtime.md) for using cached results in task environments. ## Table of Contents @@ -21,7 +21,7 @@ import torch import numpy as np from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.robots import DexforceW1Cfg -from embodichain.lab.sim.workspace import ( +from embodichain.lab.sim.motion.workspace import ( WorkspaceAnalyzer, WorkspaceAnalyzerConfig, AnalysisMode @@ -122,7 +122,7 @@ print(f"Reachable points: {results['num_reachable']}/{results['num_samples']}") ### Basic Configuration ```python -from embodichain.lab.sim.workspace.configs import VisualizationConfig +from embodichain.lab.sim.motion.workspace.configs import VisualizationConfig # Basic configuration with defaults config = WorkspaceAnalyzerConfig( @@ -164,10 +164,10 @@ import torch import numpy as np from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.robots import DexforceW1Cfg -from embodichain.lab.sim.workspace import ( +from embodichain.lab.sim.motion.workspace import ( WorkspaceAnalyzer, WorkspaceAnalyzerConfig, AnalysisMode ) -from embodichain.lab.sim.workspace.configs import VisualizationConfig +from embodichain.lab.sim.motion.workspace.configs import VisualizationConfig # Setup simulation sim = SimulationManager(SimulationManagerCfg(headless=False, sim_device="cpu")) diff --git a/docs/source/guides/add_robot.rst b/docs/source/guides/add_robot.rst index 3601c37d4..7227b81fb 100644 --- a/docs/source/guides/add_robot.rst +++ b/docs/source/guides/add_robot.rst @@ -104,5 +104,5 @@ See Also -------- - :doc:`/tutorial/robot` — Using robots in simulation -- :doc:`/overview/sim/solvers/index` — IK solver reference +- :doc:`/overview/sim/motion/solvers/index` — IK solver reference - :doc:`/resources/robot/index` — Existing robot documentation diff --git a/docs/source/index.rst b/docs/source/index.rst index b1a80b366..e886e7146 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -37,6 +37,7 @@ Table of Contents overview/sim/index overview/task_program/index overview/gym/index + overview/trajectory_generation overview/rl/index .. toctree:: diff --git a/docs/source/overview/gym/index.rst b/docs/source/overview/gym/index.rst index 439445cfd..48db66eca 100644 --- a/docs/source/overview/gym/index.rst +++ b/docs/source/overview/gym/index.rst @@ -196,6 +196,8 @@ Choosing Where to Start - Use :doc:`reward_functors` when composing RL reward terms. - Use :doc:`dataset_functors` when recording demonstrations or exporting datasets. Use this page for structured episode data, not debug video capture. +- Use :doc:`/overview/trajectory_generation` when repeating candidate rollouts + from a caller-provided fixed initial state with exclusive batch ownership. Documentation Quality Notes --------------------------- diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index e7a21f3dc..298ff48b1 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -926,7 +926,7 @@ See {doc}`builtin_actions` for the shipped skill catalog and visual demos, and robot profiles, and effects - {doc}`../../task_program/robot_profiles` — embodiment resources and presets - {doc}`../../task_program/index` — semantic task compilation and execution -- {doc}`../planners/motion_generator` — the motion generator owned by the engine +- {doc}`../motion/planners/motion_generator` — the motion generator owned by the engine - {doc}`../sim_robot` — robot control parts and kinematic configuration - {doc}`/tutorial/atomic_actions` — static, closed-loop, and recovery examples - `scripts/tutorials/atomic_action/moving_target_recovery.py` — runnable runner diff --git a/docs/source/overview/sim/index.rst b/docs/source/overview/sim/index.rst index 17607bcbc..650d6e0ef 100644 --- a/docs/source/overview/sim/index.rst +++ b/docs/source/overview/sim/index.rst @@ -38,13 +38,11 @@ The simulation stack can be read from the bottom up: | `-- contact sensors |-- browser visualization | `-- Viser scene, cameras, overlays, and runtime telemetry - |-- solvers - | |-- forward kinematics - | |-- inverse kinematics - | `-- differential kinematics - |-- planners - | |-- joint-space and Cartesian trajectory generation - | `-- time parameterization and sampling utilities + |-- motion + | |-- solvers: forward, inverse, and differential kinematics + | |-- planners: paths, time parameterization, and sampling + | |-- workspace: reachability analysis and runtime queries + | `-- trajectory_augmentation: candidates, coverage, and generation accounting `-- atomic actions `-- reusable manipulation primitives built from assets, solvers, and planners @@ -87,15 +85,22 @@ Submodule Relationships - Produce perception data such as color, depth, segmentation, stereo disparity, and contacts. - Attach to world frames, robot links, or monitored bodies and return batched tensors for downstream policies or datasets. - * - Solvers + * - Motion / solvers - Compute FK, IK, and differential kinematics for robots and articulations. - Translate task-space goals into joint-space commands used by planners, controllers, and actions. - * - Planners + * - Motion / planners - Generate joint-space or Cartesian trajectories with interpolation, timing, and feasibility handling. - Use robot state and solver results to produce trajectories that can be replayed in the manager loop. + * - Motion / workspace + - Analyze reachability and cache workspace samples for runtime queries. + - Supplies candidate robot configurations and poses for motion workflows. + * - Motion / trajectory augmentation + - Vary annotated trajectories and track candidate budgets and coverage. + - Receives planning, rollout, validation, and persistence results from + explicit host integrations. * - Atomic actions - Package complete manipulation primitives such as move, pick, and place. - Compose semantic targets, solvers, planners, and robot control into @@ -134,9 +139,10 @@ Choosing Where to Start - Use :doc:`sim_sensor` when adding camera, stereo, or contact observations. - Use :doc:`viser_visualization` when inspecting a headless or remote scene in a browser. -- Use :doc:`solvers/index` when a robot needs FK, IK, or velocity-level +- Use :doc:`motion/index` for the shared robot motion package and its boundaries. +- Use :doc:`motion/solvers/index` when a robot needs FK, IK, or velocity-level kinematics. -- Use :doc:`planners/index` when a target pose or joint goal must become a +- Use :doc:`motion/planners/index` when a target pose or joint goal must become a time-ordered trajectory. - Use :doc:`atomic actions ` when building scripted manipulation from reusable motion primitives. @@ -169,6 +175,5 @@ See Also sim_assets.md sim_sensor.md viser_visualization.md - solvers/index - planners/index + motion/index atomic_actions/index diff --git a/docs/source/overview/sim/motion/index.rst b/docs/source/overview/sim/motion/index.rst new file mode 100644 index 000000000..8eaa6c471 --- /dev/null +++ b/docs/source/overview/sim/motion/index.rst @@ -0,0 +1,47 @@ +Robot Motion +============ + +The :mod:`embodichain.lab.sim.motion` package groups four related capabilities: +kinematic solvers, trajectory planners, workspace analysis, and expert trajectory +augmentation. Simulation objects and atomic actions use these capabilities to +translate robot goals into executable motion. + +Choose the owning subpackage when importing an API. The ``motion`` parent loads +its children on access and does not re-export their classes and functions; +public imports still use the normal ``lab`` and ``sim`` initialization path. + +.. list-table:: + :header-rows: 1 + :widths: 24 76 + + * - Subpackage + - Responsibility + * - ``motion.solvers`` + - Forward, inverse, and differential kinematics. + * - ``motion.planners`` + - Joint-space and Cartesian paths, collision-aware planning, and time + parameterization. + * - ``motion.workspace`` + - Offline reachability analysis, workspace caches, and runtime sampling. + * - ``motion.trajectory_augmentation`` + - Explicit qpos templates and candidates, constrained geometric and timing + variation, measured coverage, and bounded generation accounting. + +The augmentation package currently provides the core contracts and operators. +The :doc:`generation host layer ` supplies +full-batch initial-state preparation, restoration, qpos execution, and a +synchronous episode sink. Its handwritten qpos runner collects explicitly free +motion using measured collision/task evidence and persistence confirmations. +Contact, held-object, and atomic-source collection remain outside this path. See the +:doc:`augmentation API ` +for the implemented boundaries. + +.. toctree:: + :maxdepth: 1 + + solvers/index + planners/index + +Workspace workflows are documented in +:doc:`/features/workspace_analyzer/index`. The complete public package surface is +listed in the :doc:`motion API `. diff --git a/docs/source/overview/sim/planners/curobo_planner.md b/docs/source/overview/sim/motion/planners/curobo_planner.md similarity index 94% rename from docs/source/overview/sim/planners/curobo_planner.md rename to docs/source/overview/sim/motion/planners/curobo_planner.md index 884f15f2a..6e8aafc44 100644 --- a/docs/source/overview/sim/planners/curobo_planner.md +++ b/docs/source/overview/sim/motion/planners/curobo_planner.md @@ -68,11 +68,11 @@ the fingers in the planned control part. A mismatch means cuRobo validates a different collision geometry from the one replayed in DexSim. Assuming the scene has been registered as shown in -{doc}`../../task_program/scene_registry`, construct the planner world from +{doc}`../../../task_program/scene_registry`, construct the planner world from that catalog: ~~~python -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( CuroboPlannerCfg, CuroboWorldCfg, MotionGenCfg, @@ -203,7 +203,7 @@ in `rigid_objects`. A sequence of objects is retained only as an advanced direct-core path; it derives names from each `uid` or an `obstacle_` fallback. Do not use that form for a registry-backed world. -The {doc}`../../task_program/scene_registry` integration performs two +The {doc}`../../../task_program/scene_registry` integration performs two higher-level checks before execution. First, all registry `STATIC ∪ DYNAMIC` IDs must exactly equal `MotionGenerator.collision_world_entity_ids`. Second, registry, derived scene @@ -296,13 +296,20 @@ robot's URDF and solver, so nothing robot-specific needs to be hardcoded: The generated YAML is cached on disk (default `$XDG_CACHE_HOME/embodichain_curobo` or `~/.cache/embodichain_curobo`) keyed by the URDF path, URDF content, control -part, tool frame, and fit parameters, so editing the URDF or changing the fit -settings regenerates automatically and subsequent inits reuse the cache. Tune the +part, tool frame, non-control initial joint values, and fit parameters. The +non-control values determine the generated `lock_joints` geometry, including the +gripper opening. Editing those values, the URDF, or fit settings regenerates the +corresponding cached profile; subsequent matching inits reuse it. Tune the fit with `CuroboPlannerCfg.auto_gen` (`fit_type="voxel"` by default for fast first-generation; `"morphit"` for best quality; `force=True` to bypass the cache). The default `sphere_density=0.1` keeps the per-link sphere count low (~80 for a Panda) so planning stays fast; raise it for tighter collision coverage. +If non-control initial joint values change while a backend is active, cached +backend reuse is rejected. Call `planner.close()` before rebuilding against the +new locked-joint configuration. Changing only controlled-joint initial values +does not invalidate the locked collision model. + ## Generate a motion MotionGenerator passes start_qpos and control_part to the cuRobo backend. For @@ -317,7 +324,7 @@ trajectory duration). ~~~python import torch -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( CuroboPlanOptions, MotionGenOptions, PlanState, @@ -374,10 +381,10 @@ After installing cuRobo V2 and configuring a CUDA simulation environment, run the Panda obstacle-avoidance demo from the repository root: ~~~bash -python examples/sim/planners/curobo_planner.py --headless --hold-steps 1 --step-repeat 1 +python examples/sim/motion/planners/curobo_planner.py --headless --hold-steps 1 --step-repeat 1 # CPU physics with CUDA planning -python examples/sim/planners/curobo_planner.py --headless --sim-device cpu +python examples/sim/motion/planners/curobo_planner.py --headless --sim-device cpu ~~~ The demo exports the DexSim `demo_block` into the cuRobo collision world via diff --git a/docs/source/overview/sim/planners/index.rst b/docs/source/overview/sim/motion/planners/index.rst similarity index 100% rename from docs/source/overview/sim/planners/index.rst rename to docs/source/overview/sim/motion/planners/index.rst diff --git a/docs/source/overview/sim/planners/motion_generator.md b/docs/source/overview/sim/motion/planners/motion_generator.md similarity index 93% rename from docs/source/overview/sim/planners/motion_generator.md rename to docs/source/overview/sim/motion/planners/motion_generator.md index 08f71b6fb..db2c6b49f 100644 --- a/docs/source/overview/sim/planners/motion_generator.md +++ b/docs/source/overview/sim/motion/planners/motion_generator.md @@ -54,12 +54,12 @@ from embodichain.lab.sim.cfg import ( JointDrivePropertiesCfg, ) -from embodichain.lab.sim.planners.motion_generator import MotionGenerator, MotionGenCfg, ToppraPlannerCfg -from embodichain.lab.sim.planners.toppra_planner import ToppraPlanOptions +from embodichain.lab.sim.motion.planners.motion_generator import MotionGenerator, MotionGenCfg, ToppraPlannerCfg +from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanOptions from embodichain.lab.sim.objects.robot import Robot -from embodichain.lab.sim.solvers.pink_solver import PinkSolverCfg -from embodichain.lab.sim.planners.utils import TrajectorySampleMethod, PlanState, MoveType -from embodichain.lab.sim.planners.motion_generator import MotionGenOptions +from embodichain.lab.sim.motion.solvers.pink_solver import PinkSolverCfg +from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod, PlanState, MoveType +from embodichain.lab.sim.motion.planners.motion_generator import MotionGenOptions # Configure the simulation sim_cfg = SimulationManagerCfg( @@ -238,5 +238,5 @@ print(f"Estimated sample count: {sample_count}") * For Cartesian interpolation, inverse kinematics (IK) is used to compute joint configurations for each interpolated pose. * Backends declare whether pre-interpolation is safe and whether their returned samples must be preserved. cuRobo V2 disables EmbodiChain Cartesian pre-interpolation and (by default) is resampled to `MotionGenOptions.sample_count`; set `CuroboPlannerCfg.preserve_plan_samples=True` to keep its raw collision-checked samples. * CuroboPlanner is optional and requires CUDA plus a matching cuRobo V2 installation; see [the cuRobo planner page](curobo_planner.md) and [NVIDIA's installation guide](https://nvlabs.github.io/curobo/latest/getting-started/installation.html). -* Run the collision-aware Panda demo with `python examples/sim/planners/curobo_planner.py --headless --hold-steps 1 --step-repeat 1`. +* Run the collision-aware Panda demo with `python examples/sim/motion/planners/curobo_planner.py --headless --hold-steps 1 --step-repeat 1`. * The sample count estimation is useful for predicting computational load and memory requirements. diff --git a/docs/source/overview/sim/planners/neural_planner.md b/docs/source/overview/sim/motion/planners/neural_planner.md similarity index 87% rename from docs/source/overview/sim/planners/neural_planner.md rename to docs/source/overview/sim/motion/planners/neural_planner.md index 632e2c488..08d9466dc 100644 --- a/docs/source/overview/sim/planners/neural_planner.md +++ b/docs/source/overview/sim/motion/planners/neural_planner.md @@ -18,7 +18,7 @@ Pre-trained checkpoints are hosted on HuggingFace and can be downloaded with ```python from embodichain.data.assets.planner_assets import download_neural_planner_checkpoint -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( MotionGenCfg, MotionGenOptions, MotionGenerator, @@ -26,7 +26,7 @@ from embodichain.lab.sim.planners import ( NeuralPlannerCfg, PlanState, ) -from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions +from embodichain.lab.sim.motion.planners.neural_planner import NeuralPlanOptions checkpoint_path = download_neural_planner_checkpoint() @@ -57,7 +57,7 @@ result = motion_generator.generate( ## Example ```bash -python examples/sim/planners/neural_planner.py --headless --device cuda +python examples/sim/motion/planners/neural_planner.py --headless --device cuda ``` The example downloads the checkpoint automatically on first run. diff --git a/docs/source/overview/sim/planners/toppra_planner.md b/docs/source/overview/sim/motion/planners/toppra_planner.md similarity index 90% rename from docs/source/overview/sim/planners/toppra_planner.md rename to docs/source/overview/sim/motion/planners/toppra_planner.md index ca02433b6..cba8a98a4 100644 --- a/docs/source/overview/sim/planners/toppra_planner.md +++ b/docs/source/overview/sim/motion/planners/toppra_planner.md @@ -14,7 +14,7 @@ ### Initialization ```python -from embodichain.lab.sim.planners.toppra_planner import ToppraPlanner, ToppraPlannerCfg +from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanner, ToppraPlannerCfg # Configuration - constraints are now specified in ToppraPlanOptions, not here cfg = ToppraPlannerCfg( @@ -26,8 +26,8 @@ planner = ToppraPlanner(cfg=cfg) ### Planning ```python -from embodichain.lab.sim.planners.utils import TrajectorySampleMethod, PlanState -from embodichain.lab.sim.planners.toppra_planner import ToppraPlanOptions +from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod, PlanState +from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanOptions # Create options with constraints and sampling parameters options = ToppraPlanOptions( diff --git a/docs/source/overview/sim/planners/trajectory_sample_method.md b/docs/source/overview/sim/motion/planners/trajectory_sample_method.md similarity index 100% rename from docs/source/overview/sim/planners/trajectory_sample_method.md rename to docs/source/overview/sim/motion/planners/trajectory_sample_method.md diff --git a/docs/source/overview/sim/solvers/differential_solver.md b/docs/source/overview/sim/motion/solvers/differential_solver.md similarity index 95% rename from docs/source/overview/sim/solvers/differential_solver.md rename to docs/source/overview/sim/motion/solvers/differential_solver.md index 5f68cad69..a04535180 100644 --- a/docs/source/overview/sim/solvers/differential_solver.md +++ b/docs/source/overview/sim/motion/solvers/differential_solver.md @@ -13,8 +13,8 @@ The `DifferentialSolver` is a differential inverse kinematics (IK) controller de ```python from embodichain.data import get_data_path -from embodichain.lab.sim.solvers.differential_solver import DifferentialSolver -from embodichain.lab.sim.solvers.differential_solver import DifferentialSolverCfg +from embodichain.lab.sim.motion.solvers.differential_solver import DifferentialSolver +from embodichain.lab.sim.motion.solvers.differential_solver import DifferentialSolverCfg cfg = DifferentialSolverCfg( urdf_path=get_data_path("UniversalRobots/UR5/UR5.urdf"), diff --git a/docs/source/overview/sim/solvers/index.rst b/docs/source/overview/sim/motion/solvers/index.rst similarity index 100% rename from docs/source/overview/sim/solvers/index.rst rename to docs/source/overview/sim/motion/solvers/index.rst diff --git a/docs/source/overview/sim/solvers/neural_ik_solver.md b/docs/source/overview/sim/motion/solvers/neural_ik_solver.md similarity index 96% rename from docs/source/overview/sim/solvers/neural_ik_solver.md rename to docs/source/overview/sim/motion/solvers/neural_ik_solver.md index f2690d016..f6681c972 100644 --- a/docs/source/overview/sim/solvers/neural_ik_solver.md +++ b/docs/source/overview/sim/motion/solvers/neural_ik_solver.md @@ -28,7 +28,7 @@ checkpoints are hosted on HuggingFace and can be downloaded with ```python from embodichain.data.assets.solver_assets import download_neural_ik_checkpoint -from embodichain.lab.sim.solvers.neural_ik_solver import NeuralIKSolverCfg +from embodichain.lab.sim.motion.solvers.neural_ik_solver import NeuralIKSolverCfg checkpoint_path = download_neural_ik_checkpoint() diff --git a/docs/source/overview/sim/solvers/opw_solver.md b/docs/source/overview/sim/motion/solvers/opw_solver.md similarity index 98% rename from docs/source/overview/sim/solvers/opw_solver.md rename to docs/source/overview/sim/motion/solvers/opw_solver.md index 8f7bff079..8599c9e96 100644 --- a/docs/source/overview/sim/solvers/opw_solver.md +++ b/docs/source/overview/sim/motion/solvers/opw_solver.md @@ -20,7 +20,7 @@ The solver is configured using the `OPWSolverCfg` class, which defines OPW param import torch import numpy as np from embodichain.data import get_data_path -from embodichain.lab.sim.solvers.opw_solver import OPWSolver, OPWSolverCfg +from embodichain.lab.sim.motion.solvers.opw_solver import OPWSolver, OPWSolverCfg cfg = OPWSolverCfg( urdf_path=get_data_path("CobotMagicArm/CobotMagicNoGripper.urdf"), diff --git a/docs/source/overview/sim/solvers/pink_solver.md b/docs/source/overview/sim/motion/solvers/pink_solver.md similarity index 96% rename from docs/source/overview/sim/solvers/pink_solver.md rename to docs/source/overview/sim/motion/solvers/pink_solver.md index bc44c9ae1..5188d57e2 100644 --- a/docs/source/overview/sim/solvers/pink_solver.md +++ b/docs/source/overview/sim/motion/solvers/pink_solver.md @@ -18,8 +18,8 @@ ```python from embodichain.data import get_data_path -from embodichain.lab.sim.solvers.pink_solver import PinkSolverCfg -from embodichain.lab.sim.solvers.null_space_posture_task import NullSpacePostureTask +from embodichain.lab.sim.motion.solvers.pink_solver import PinkSolverCfg +from embodichain.lab.sim.motion.solvers.null_space_posture_task import NullSpacePostureTask posture_task = NullSpacePostureTask( cost=1e-3, diff --git a/docs/source/overview/sim/solvers/pinocchio_solver.md b/docs/source/overview/sim/motion/solvers/pinocchio_solver.md similarity index 95% rename from docs/source/overview/sim/solvers/pinocchio_solver.md rename to docs/source/overview/sim/motion/solvers/pinocchio_solver.md index fa2f0dec9..9524828aa 100644 --- a/docs/source/overview/sim/solvers/pinocchio_solver.md +++ b/docs/source/overview/sim/motion/solvers/pinocchio_solver.md @@ -15,8 +15,8 @@ The `PinocchioSolver` is a high-precision inverse kinematics (IK) solver for rob ```python from embodichain.data import get_data_path -from embodichain.lab.sim.solvers.pinocchio_solver import PinocchioSolver -from embodichain.lab.sim.solvers.pinocchio_solver import PinocchioSolverCfg +from embodichain.lab.sim.motion.solvers.pinocchio_solver import PinocchioSolver +from embodichain.lab.sim.motion.solvers.pinocchio_solver import PinocchioSolverCfg cfg = PinocchioSolverCfg( urdf_path=get_data_path("UniversalRobots/UR5/UR5.urdf"), diff --git a/docs/source/overview/sim/solvers/pytorch_solver.md b/docs/source/overview/sim/motion/solvers/pytorch_solver.md similarity index 96% rename from docs/source/overview/sim/solvers/pytorch_solver.md rename to docs/source/overview/sim/motion/solvers/pytorch_solver.md index 7677a3cb7..cc0946351 100644 --- a/docs/source/overview/sim/solvers/pytorch_solver.md +++ b/docs/source/overview/sim/motion/solvers/pytorch_solver.md @@ -18,8 +18,8 @@ The solver is configured using the `PytorchSolverCfg` class, which allows detail ```python from embodichain.data import get_data_path -from embodichain.lab.sim.solvers.pytorch_solver import PytorchSolver -from embodichain.lab.sim.solvers.pytorch_solver import PytorchSolverCfg +from embodichain.lab.sim.motion.solvers.pytorch_solver import PytorchSolver +from embodichain.lab.sim.motion.solvers.pytorch_solver import PytorchSolverCfg import torch cfg = PytorchSolverCfg( diff --git a/docs/source/overview/sim/solvers/srs_solver.md b/docs/source/overview/sim/motion/solvers/srs_solver.md similarity index 98% rename from docs/source/overview/sim/solvers/srs_solver.md rename to docs/source/overview/sim/motion/solvers/srs_solver.md index 8e2d73563..94374e73f 100644 --- a/docs/source/overview/sim/solvers/srs_solver.md +++ b/docs/source/overview/sim/motion/solvers/srs_solver.md @@ -36,7 +36,7 @@ from embodichain.lab.sim.robots.dexforce_w1.types import ( from embodichain.lab.sim.robots.dexforce_w1.params import ( W1ArmKineParams, ) -from embodichain.lab.sim.solvers.srs_solver import SRSSolver, SRSSolverCfg +from embodichain.lab.sim.motion.solvers.srs_solver import SRSSolver, SRSSolverCfg arm_params = W1ArmKineParams( arm_side=DexforceW1ArmSide.RIGHT, diff --git a/docs/source/overview/sim/solvers/ur_solver.md b/docs/source/overview/sim/motion/solvers/ur_solver.md similarity index 99% rename from docs/source/overview/sim/solvers/ur_solver.md rename to docs/source/overview/sim/motion/solvers/ur_solver.md index 57a5524f5..73a594f21 100644 --- a/docs/source/overview/sim/solvers/ur_solver.md +++ b/docs/source/overview/sim/motion/solvers/ur_solver.md @@ -49,7 +49,7 @@ default to sensible values and rarely need to be overridden. ```python import torch -from embodichain.lab.sim.solvers.ur_solver import URSolver, URSolverCfg +from embodichain.lab.sim.motion.solvers.ur_solver import URSolver, URSolverCfg cfg = URSolverCfg( ur_type="ur10", # one of: ur3, ur3e, ur5, ur5e, ur10, ur10e diff --git a/docs/source/overview/sim/sim_robot.md b/docs/source/overview/sim/sim_robot.md index 2e8e36261..80c9f842b 100644 --- a/docs/source/overview/sim/sim_robot.md +++ b/docs/source/overview/sim/sim_robot.md @@ -22,7 +22,7 @@ A `Robot` must be spawned within a `SimulationManager`. import torch from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.objects import Robot, RobotCfg -from embodichain.lab.sim.solvers import SolverCfg +from embodichain.lab.sim.motion.solvers import SolverCfg # 1. Initialize Simulation Environment # Note: Use 'sim_device' to specify device (e.g., "cuda:0" or "cpu") diff --git a/docs/source/overview/task_program/scene_registry.md b/docs/source/overview/task_program/scene_registry.md index 4d2c708ba..13cc81f14 100644 --- a/docs/source/overview/task_program/scene_registry.md +++ b/docs/source/overview/task_program/scene_registry.md @@ -275,7 +275,7 @@ cuRobo world, derive both the explicit `registry_id -> RigidObject` mapping and the dynamic-obstacle ID list from the registry: ```python -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( CuroboPlannerCfg, CuroboWorldCfg, MotionGenCfg, @@ -383,5 +383,5 @@ provide alias normalization or registry/provider/planner construction checks. See {doc}`index` for manifest and semantic-call integration, {doc}`../sim/atomic_actions/index` for snapshot grounding and recovery -semantics, and {doc}`../sim/planners/curobo_planner` for cuRobo world +semantics, and {doc}`../sim/motion/planners/curobo_planner` for cuRobo world representation and frame details. diff --git a/docs/source/overview/trajectory_generation.rst b/docs/source/overview/trajectory_generation.rst new file mode 100644 index 000000000..5f10b01e6 --- /dev/null +++ b/docs/source/overview/trajectory_generation.rst @@ -0,0 +1,566 @@ +Fixed-Scene Trajectory Generation +================================== + +Repeated expert rollouts need the same declared scene and initial state before +each candidate executes. The :mod:`embodichain.lab.trajectory_generation` layer +provides that preparation lifecycle for an entire simulator batch, either +directly or through an ``EmbodiedEnv``. Its synchronous qpos runner connects +candidate planning, physical execution, measured validation, and +persistence of accepted episodes. It consumes a scene that the caller already +created; it does not sample layouts. + +The :mod:`embodichain.lab.sim.motion.trajectory_augmentation` package provides +qpos candidates, constrained variation, coverage, and generation accounting. +Two collection paths are available: free motion through the existing locked-joint +planner, and a bounded CPU-physics PickUp integration with full-joint geometry +and physical contact validation. PickUp consumes offline atomic compilations; +AtomicActionRuntime recovery/feedback execution, contact-aware Gym collection, +and the unified generation CLI remain unimplemented. Initial-state preparation +alone does not certify a candidate trajectory or confirm that an episode was saved. + +Supported Initial States +------------------------ + +The physical adapter supports one configuration-owned fixed-base robot and all +ordinary rigid objects in the simulator. Every entity must cover every simulator +row. Cases are supplied in physical row order, so one batch can retain different +fixed cases in its rows. Partial-row restoration and replica rebinding are not +implemented. + +Snapshots own local poses, rigid velocities, complete joint positions and +velocities, drive targets, and joint effort state. They include mimic and gripper +joints. Additional articulations, object groups, deformables, rigid constraints, +and robots using USD articulation properties are rejected. The initial state +must not depend on pending external forces or an active contact constraint; +this is episode initialization, not a mid-contact physics checkpoint. + +The physical adapter has been exercised against a real headless CPU backend +with physics both enabled and disabled, using one robot and one cube. Those +checks cover backend state reads, writes, and restoration. The PickUp collection +example below additionally exercises repeated restoration after real grasps; +it does not restore a contact solver checkpoint. + +A Trusted Preparation Profile +----------------------------- + +The caller supplies an +:class:`~embodichain.lab.trajectory_generation.initial_state.InitialStateProfile` +implemented for the task, controller, and scene configuration: + +* ``prepare`` deterministically initializes task state and controller history, + including state outside the standard Gym episode managers. It must preserve + the fixed scene conditions. +* ``signature`` identifies fixed geometry, physics, materials, visuals, sensors, + and control conditions. It excludes evolving poses and joint positions and + must change when a condition relevant to the case changes. +* ``verify`` returns nonempty, passing ``ValidationResult`` evidence for every + declared case after settling. It must check the task-specific initial + conditions; an unconditional passing result is not a preparation policy. +* ``physics_dt`` declares the settling clock and must match ``env.physics_dt`` + for Gym. ``settling_steps`` controls preparation-only simulation steps. +* ``allowed_interval_events`` names Gym interval events explicitly checked to + preserve the fixed conditions. Any active interval event outside this list + prevents host acquisition or reuse. + +Callbacks come from trusted integration code. There is no generic safe default +profile and no loading of profile callbacks from job YAML strings. The physical +adapter checks scene structure and numeric state; it cannot infer the fixed +condition identity or reset controller-specific memory for the profile. + +Preparation and Rollout Order +----------------------------- + +:class:`~embodichain.lab.trajectory_generation.initial_state.FixedSceneHost` +owns the complete simulator batch. With Gym it also acquires an environment +generation lease, which suppresses automatic resets and rejects ordinary +``env.reset()`` until ownership is released. The owner must serialize all host +access and remain the sole physical stepper. Direct simulation hosts and Gym +leases share ownership of the same simulator batch; a second owner through +either interface is rejected. + +The first ``acquire_case(cases)`` performs deterministic preparation, settling, +and verification, then captures the prepared physical states. Later +``restore_initial()`` calls reconstruct those captured states. In Gym, each +preparation follows this order: + +1. Invalidate the previous epoch and active Task Program bridge; disable + stepping while preparation is incomplete. +2. Discard pending camera and episode recording state. The caller must already + own copies of the previous episode and its validation evidence. +3. Run deterministic profile preparation, restore the captured physical state + when one exists, then perform the explicit settling steps. +4. Reset observation history, reward, and dataset managers, then verify task + initial conditions and, for restoration, the captured physical fields. The + validator observes the new episode's manager state. A failed attempt publishes + no execution binding. +5. Obtain the initial observation and seed enabled recorders from the verified + state. + +Preparation does not run reset/interval events, reseed the environment RNG, or +call normal Gym ``step``. Profile callbacks must not call ``env.step`` or +``env.reset``. Settling and the articulation root setter may advance physics; +those preparation updates are outside the recorded rollout. + +``PreparedBatch`` is an execution binding for one host and epoch. Check it with +``host.assert_current(binding)`` before issuing candidate commands. Restoration, +failed preparation, and ownership release invalidate previous bindings. A new +binding requires new runtime requests and Task Program bridges where applicable. + +``host.snapshots(binding)`` returns independent planning snapshots of the +captured initial batch. For Gym, ``host.initial_observation(binding)`` copies the +first observation already produced by preparation, preserving observation-history +semantics without another ``get_obs`` call. It returns ``None`` for direct sim. + +Integrating an Existing Gym Scene +--------------------------------- + +The following helper assumes an already-created unwrapped ``EmbodiedEnv``, one +``SceneCase`` per row, a task-specific trusted profile, and explicit execution +and evidence-copy callbacks. ``execute_candidate`` uses the normal environment +step path; ``freeze_episode`` copies all data needed by later validation and +persistence before the next restoration discards the live buffers. + +.. code-block:: python + + from embodichain.lab.trajectory_generation.initial_state import FixedSceneHost + from embodichain.lab.trajectory_generation.integrations.sim import ( + SimInitialStateAdapter, + ) + + def run_two_candidates(env, cases, profile, execute_candidate, freeze_episode): + adapter = SimInitialStateAdapter(env.sim, env.robot) + episodes = [] + with FixedSceneHost(adapter, profile, env=env) as host: + binding = host.acquire_case(cases) + for candidate_index in range(2): + if candidate_index: + binding = host.restore_initial() + host.assert_current(binding) + execute_candidate(env, binding, candidate_index) + episodes.append(freeze_episode(env, binding)) + return episodes + +For direct simulation, construct the same host without ``env=``. The caller +then owns observation refresh, rollout timing, recording, and task execution. +Closing the host releases ownership and invalidates bindings without implicitly +saving or resetting the final episode. + +Checking Free-Motion Candidates +------------------------------- + +The +:class:`~embodichain.lab.trajectory_generation.integrations.planning.EnvRowMotionPlanner` +adapter connects explicit EEF or qpos candidates to a configured +``MotionGenerator`` and the actual robot batch. Logical candidate count ``C`` can +exceed physical row count ``B``: candidates are checked in rounds, +with at most one candidate assigned to its source row in each round. Snapshots +must cover the complete batch, and backend calls retain size ``B``. + +For EEF input, +:class:`~embodichain.lab.trajectory_generation.integrations.planning.EEFPath` +declares local-arena TCP samples, a source row, identity, explicit arrival +intervals, and phases. ``plan_eef`` propagates successful IK seeds between +samples and checks FK position/orientation residuals. Supplying solved joint +samples preserves that branch without another IK solve. The output is a +full-joint ``CandidateTrajectoryBatch`` and one validation result per candidate. + +For qpos input, ``validate_qpos`` checks the original candidate paths. It inserts +bounded joint-space samples between every adjacent pair, including phase +boundaries, before asking the backend to check self/environment collision. +Sampling changes only the check inputs. Original candidate positions and timing +remain intact. + +This initial adapter supports fully annotated free motion, explicitly declared +empty ``held_object_ids``, and unchanged uncontrolled joints at the configured +initial values. Contact/hold phases, unannotated paths, held-object sweeps, and +changing locked joints +such as gripper closure are outside this capability. Missing collision support +or a path requiring more than ``max_validation_samples`` also returns +``unavailable``. Such candidates cannot enter accepted rollout generation. + +An unchanged gripper command does not guarantee unchanged measured finger +positions. The measured path must also match the backend's locked-joint values +within ``1e-6``. A real Panda run under ordinary gravity exceeded this bound +through finger drift and was rejected with no committed episode, even though its +task, motion-quality, and speed/acceleration checks passed. This is an unsupported +locked-model state for measured collision validation. + +The host must keep robot roots and canonical-ID collision obstacles synchronized +with the snapshots. A passing sampled collision result does not provide a +continuous collision guarantee, controller speed/acceleration validation, or +task success. Those checks remain explicit later gates; this adapter is not yet +a contact-aware PickUp planner. + +The adapter's real cuRobo smoke uses one Panda with CPU physics and CUDA collision +checking. A short free path passes with a distant dynamic cube and fails after +the cube is moved to the TCP, reusing the same backend. This checks dynamic-world +collision updates for that sampled path; held-object and contact semantics remain +outside the supported scope. + +Executing Qpos Candidates +------------------------- + +:class:`~embodichain.lab.trajectory_generation.execution.QposRolloutExecutor` +accepts a prepared host binding, one single-row candidate or ``None`` per +physical row, and stable episode/commit IDs. Each candidate contains an initial +sample followed by commands on the host's fixed control clock. Inactive and +mimic joints must remain unchanged; the executor checks initial state, joint +limits, layout, and payload size before submitting a command. + +The executor's byte limit includes tensors plus a 64 KiB metadata reserve. +Executor metadata is limited to 32 KiB so the runner and sink can add their +evidence within that reserve. Initial observation size and later schema changes +are checked before the executor copies observation payloads. + +Gym execution uses the existing demo loop and ``ControllerAction``. It reads the +actual controller position targets before physics and consumes the existing +``env.step`` observation afterward. Direct sim uses explicit integer physics +substeps and a full-batch observation callback. A Gym encoder can transform the +prepared/step observations; its default flattens all tensor fields without +dropping channels. The executor always includes measured full-joint +``joint_positions``. +The default encoder requires nonempty string keys without embedded dots and +tensor leaves; other observation structures require an explicit encoder. + +Timestamps are differences of ``SimulationManager.simulation_time`` and each +transition must advance exactly one control period within tolerance. That clock +counts manager-owned updates; external direct backend updates are outside it and +must not occur during a rollout owned by this executor. + +Completed rows freeze their causal frames while other rows finish. Idle, +finished, or terminated rows receive holds, and those holds do not enter the +returned training sequence. Installing holds clears velocity/effort targets +without advancing physics. A row with no complete transition returns ``None``; +an interrupted row with some complete transitions returns rejected evidence. +``executor.last_failures`` keeps bounded candidate-specific runtime reasons even +when no complete transition exists. It is read-only and resets at the next call; +the runner records these reasons when releasing failed candidates. +The final task validator consumes owned row observations, actions, and times. +It must implement the task's real success conditions and include its configured +check ID. + +The executor checks root and rigid-object pose stability at every observation +boundary and adds ``fixed_collision_world`` evidence. The runner uses that +evidence when rechecking the measured qpos path against the initial collision +snapshot. Boundary checks and sampled collision checks do not establish +continuous collision freedom or support contact/held-object tasks. + +Assembling the Handwritten Runner +--------------------------------- + +:class:`~embodichain.lab.trajectory_generation.runner.GenerationRunner` combines +the explicit host, planning, execution, and persistence interfaces documented +here. The caller supplies qualified references, initial-state and validation +profiles, and full-joint limits through +:class:`~embodichain.lab.trajectory_generation.runner.MotionLimitsProfile`. + +Every host rigid-object physical UID must appear in the backend's collision-world +and dynamic-pose entity IDs. This lets the runner supply the captured initial +poses for all rows after preparation. Dynamic here means updateable collision +poses, including physically static bodies. A baked static world is insufficient; +different relative layouts across rows require per-environment backend worlds. +Semantic aliases are not yet connected. The trusted profile must also establish +coverage of relevant implicit planes and other geometry outside the rigid +registry. The executor's per-row byte limit must not exceed the sink or job +pending-byte limit. + +.. code-block:: python + + from embodichain.lab.trajectory_generation.runner import GenerationRunner + + # All ports are explicitly constructed for the same scene, robot, and clock. + # host owns the batch, but has not yet called acquire_case(). + runner = GenerationRunner( + cfg, host, planner, executor, sink, motion_limits=motion_limits + ) + report = runner.run(cases, templates) + +``cases`` and ``templates`` contain one item per physical row in matching order. +The runner verifies profile/source/template IDs and clock compatibility, proposes +allowed residual and retiming variants, validates their paths and motion limits, +then reserves capacity before each rollout. Measured paths, task evidence, +speed/acceleration, and path-length/duration quality gates determine which +episodes reach the sink. Write retries reuse frozen evidence. + +Planning evidence is retained as ``planned_*`` checks. The measured-path collision +check, ``actual_motion_limits``, and ``motion_quality`` remain separate gates; +planned feasibility alone cannot qualify the saved episode. + +``run`` is single-use and closes the supplied host and sink on exit. It returns +counts, coverage, audit, resolved configuration, and ``target_reached`` and writes +``generation_report.json`` under the sink root. Reaching a proposal, rollout, or +wall-time budget can finish without reaching the collection target. The PickUp +adapter below adds contact-aware offline atomic replay; a unified configuration +registry and generation CLI remain subsequent work. + +Running the Real Free-Motion Example +------------------------------------- + +From the repository root, with the simulation dependencies, CUDA, cuRobo, and +LeRobot installed, use a new or empty output directory: + +.. code-block:: bash + + python examples/sim/motion/trajectory_generation/free_motion.py --output /tmp/ur5-free-motion + +The example constructs its own fixed scene and trusted preparation/validation +profiles. It uses a pure-arm UR5 with ordinary gravity, one CPU physics row, +CUDA collision checking, and no desktop window. The default run has no camera. +A registered ground proxy +matches the simulator's implicit floor so the collision world covers that +geometry. ``--cuda-device`` selects the GPU, defaulting to ``0``. + +One second of motion produces 21 measured observations and 20 actual commands. +The verified run committed one episode after planned and measured collision, +execution, task, quality, and motion-limit checks passed. It measured about +``0.078021`` rad of movement, ``0.003212`` rad endpoint error, and ``0.010872`` rad +maximum tracking error. These are results of this short example, not a PickUp +or throughput acceptance result. + +The script prints counts, ``target_reached``, and the report path. The output +contains ``generation_report.json``, a committed-shard ``manifest.json``, and +the LeRobot dataset plus required evidence files described below. A run that +does not reach its target exits with status ``1``. Only manifest-listed shards +are committed training data. + +To save an offscreen video of a larger, five-second motion, install +``imageio-ffmpeg`` in the same environment and run: + +.. code-block:: bash + + python examples/sim/motion/trajectory_generation/free_motion.py --output /tmp/ur5-video --record-video --duration 5 --joint-displacement 0.4 + +``--record-video`` adds a 640 x 480 RGB camera and streams its actual execution +frames to ``preview.mp4`` using H.264 at 20 fps. It includes the initial frame +and each subsequent observation, including the terminal frame: five seconds of +physical motion produces 101 frames, so the video lasts 5.05 seconds. Numeric +LeRobot observations retain their existing schema; this preview is a separate +visual record. A video may also show a rejected attempt, so use +``generation_report.json`` and the committed manifest to assess whether the +episode qualified as expert data. + +``--duration`` accepts 0.05 through 30 seconds in multiples of 0.05; its default +is 1 second. ``--joint-displacement`` controls the first arm joint's positive +displacement, up to 0.5 rad, and defaults to 0.08 rad. Changing these values +preserves the validation gates and does not guarantee acceptance. + +Using ``--robot panda`` with a separate empty output directory runs the rejection +example: ordinary-gravity finger drift exceeds the locked-model tolerance, so +measured collision validation is unavailable and no episode is committed. +The real integration test covers both outcomes: + +.. code-block:: bash + + pytest tests/lab/trajectory_generation/test_runner_real.py --run-gpu -m gpu -q + +Both cases passed. The UR5 case also decodes the saved H.264 preview and checks +that frame timestamps match all measured observations, including the terminal +state, without extending the physical rollout. Real acceptance currently covers +this one-row direct-sim free-motion example; Gym lifecycle/execution contracts +have separate CPU tests. PickUp has a separate collection example below. The +full source/runtime and sim/Gym qualification matrix remains pending. + +Parallel Cube-Grasp Augmentation Preview +---------------------------------------- + +The standalone ``cube_grasp_parallel.py`` example compares grasp-pose and +transit-path augmentation with four synchronized physical rows: + +.. code-block:: bash + + python examples/sim/motion/trajectory_generation/cube_grasp_parallel.py --output /tmp/cube-grasp-parallel + +Each row has the same initial UR5/PGI configuration and the same 5 cm cube on a +bench, with normal gravity. The rows cross two object-relative grasp +orientations (0 and 90 degrees) with reference and augmented transit paths. +``rotate_grasp_about_object_axis`` uses the cube's quarter-turn symmetry to +derive TCP goals. Existing ``MoveEndEffector`` and ``PickUp`` Atomic Skills +replan each goal using the UR IK solver and ``ik_interp`` strategy. +``joint_residual`` changes only the transit before the pre-grasp corridor, +preserving its endpoints and the later approach, close, and lift commands. + +Four offscreen cameras capture the same physics tick and stream a 1280 x 1056, +20 fps four-panel ``preview.mp4``. Colored trails show measured TCP motion; +labels show grasp orientation, phase, time, and measured cube lift. The +``--seed`` option selects the local residual random stream (default ``13``), +and ``--cuda-device`` selects the renderer GPU. Simulation dependencies, the +UR5/PGI assets, TOPPRA, and ``imageio[ffmpeg]`` are required; this example does +not initialize a cuRobo collision backend. Use a new or empty output directory. + +``rollout.npz`` stores reference and commanded qpos, measured full-joint qpos, +TCP/cube poses, grasp goals, and measured timestamps. ``report.json`` records +augmentation factors, phase boundaries, path separation, and per-row outcomes. +The final hold requires at least 12 cm of sustained cube lift, at most 1 cm of +TCP-relative position drift, and a cube-to-TCP distance below 6 cm. A run exits +with status 1 if any row fails. Cubes are moved through physical gripper contact; +no attachment constraint or pose update supplies the lift. + +The default run passed all four grasp checks, with approximately 17 cm of +maximum separation between paired transit paths. Its 13.55 seconds of physical +motion produces 272 frames including the initial and terminal observations. +This validates the demonstrated grasp behavior. Contact-aware collision +certification and expert LeRobot qualification remain outside this preview; +use the separate PickUp collection example below for expert-data validation. + +Saving Accepted Episodes +------------------------ + +:class:`~embodichain.lab.trajectory_generation.sinks.LeRobotEpisodeSink` consumes +an owned ``ExpertEpisode`` whose validation checks have passed. Each submission +creates a local LeRobot dataset shard containing one episode. The sink has no +live environment dependency and can receive frozen data from either a sim or Gym +host. The caller still owns task/trajectory validation and episode construction. + +An episode with ``T`` commands is stored as: + +.. list-table:: + :header-rows: 1 + :widths: 26 74 + + * - Artifact + - Content + * - ``dataset/`` + - ``T`` LeRobot frames pairing observations ``0..T-1`` with actual commands; + RGB observations are stored as images. + * - ``terminal.npz`` + - Observation ``T`` and all ``T+1`` measured timestamps, including their + original time origin. + * - ``episode.json`` + - Candidate/episode/commit identity, action representation, validation, + phases, feature mapping, and user metadata. + * - Collection ``manifest.json`` + - The committed episode-to-shard mapping. Unlisted shards are not committed. + +The complete causal sequence consists of the LeRobot training frames plus the +required terminal sidecar. Measured timestamps must match the configured integer +``fps`` within ``timestamp_tolerance``. LeRobot frame timestamps are relative to +the episode origin; the exact measured clock remains in the sidecar. + +Numeric observations have shape ``(T+1,D)`` or ``(T+1,M,N)`` and dtype float32, +float64, int32, int64, or uint8. Matrices flatten to vectors in C order for +LeRobot, with their source layout in ``observation_shapes`` and their terminal +matrix preserved in the sidecar. RGB images use uint8 ``(T+1,H,W,3)``. Actions use float32 or +float64 ``(T,A)``. Unsupported layouts, conflicting feature names, invalid +clocks, and payloads exceeding ``max_episode_bytes`` are rejected before episode +writes. The byte limit covers raw tensors and metadata per submission. + +After writing, the sink finalizes the LeRobot writer, opens every training frame +and image again, verifies both sidecars, and replaces and reads back the manifest. +Only then does ``submit`` return ``receipt.confirmed=True``. Persistence errors +return an unconfirmed receipt with an error message; they do not count toward +confirmed generation coverage. + +.. code-block:: python + + from embodichain.lab.trajectory_generation.sinks import LeRobotEpisodeSink + + # episode already owns the accepted rollout observations and actual commands. + with LeRobotEpisodeSink("outputs/accepted-episodes", fps=50) as sink: + receipt = sink.submit(episode) + if not receipt.confirmed: + # Reuse the same episode and commit ID; distinguish this submission. + receipt = sink.submit(episode, submission_id=1) + if not receipt.confirmed: + raise RuntimeError(receipt.error) + +The output directory must be new or empty. An identical payload under a confirmed +commit ID only reads and verifies the existing artifacts; a different payload +under that ID raises. A confirmed duplicate does not rewrite files or create +another logical episode. +Readable sealed data can be reused after a metadata failure, and incomplete +uncommitted data can be rebuilt at the same shard path. This retry behavior +applies within the owning sink process; reopening an existing collection after +a process restart is not supported. + +When using ``GenerationSession``, first reserve the measured episode with +``accept_episode``; submit only when it returns ``True``. Apply every returned +receipt with ``apply_receipt``. +After a failed write, use ``retry_write(commit_id)`` to obtain the retained +episode and its next submission ID before resubmitting. Successful write retries +reuse the rollout; they do not generate another physical attempt. + +Submission is synchronous and serialized. ``drain()`` has no deferred receipts, +and ``close()`` releases ownership without another save operation. Confirmations +cover closed, readable artifacts; they do not promise restart recovery or +power-loss durability. + +Public interfaces and supported snapshot fields are documented in the +:doc:`trajectory-generation API `. + +Contact-Validated Parallel PickUp Collection +-------------------------------------------- + +``examples/sim/motion/trajectory_generation/cube_pickup_collection.py`` reuses +the cube preview's four-row scene and MoveEndEffector/PickUp compilation, then +connects them to ``FixedSceneHost``, ``GenerationRunner``, ``GenerationSession`` +and ``LeRobotEpisodeSink``. Each row receives its own local residual proposal +in transit; approach, closing, lift and hold remain protected. Two grasp goals +use the cube's 0/90-degree symmetry. All four rows share the physical clock. + +.. code-block:: bash + + python -m pip install -e '.[trajectory-generation]' 'imageio[ffmpeg]' + python examples/sim/motion/trajectory_generation/cube_pickup_collection.py \ + --output /tmp/cube-experts --episodes 8 --record-video + +The output directory must be new or empty. ``--episodes`` defaults to 8 and is +bounded to 1–64; ``--seed`` defaults to 13. ``--cuda-device`` selects the renderer. +Physics uses CPU at 200 Hz with normal gravity; control/video use 20 Hz. +The UR5 arm stiffness is 200000 with the existing damping, and the PGI open +command sits 1 mm inside its joint limit to avoid limit overshoot during transit. +These values are fixed before initial-state capture; the same endpoint/tracking +and collision tolerances apply to both grasp orientations. + +The PickUp integration adds the following mandatory evidence: + +* Planned full-joint geometry includes both fingers and their mimic coupling. + Conservative convex hulls come from URDF collision shapes, not render meshes. + Joint-segment sampling is bounded; the held cuboid follows the TCP after lift + begins. The implicit ground is represented explicitly in the checker. +* Actual geometry uses measured joints, including passive fingers, and the + measured cube pose. Only the declared target may move; robot root and other + scene objects must retain their captured poses at observation boundaries. +* Native CPU contacts are read after every 5 ms physics step. The profile allows + only declared fixed mounting contacts, cube/support contact through initial + lift-off, and finger/cube contact during closing/lift/hold. During approach, + finger contact is confined to the declared entry region; this PGI example uses + 6 cm to include fingertip overhang beyond its TCP. Other contacts, unknown + bodies, cross-row pairs, more than 2 mm penetration, and buffer-budget overflow + reject the evidence. URDF self pairs within two kinematic hops are structural + exclusions shared with the existing cuRobo policy. +* During the 1.5-second terminal hold, each finger must report positive normal + impulse in at least 95% of physics samples and the cube must remain at least + 12 cm above its initial height. Translation and rotation drift relative to TCP + are checked from lift through hold, with 1 cm / 0.15 rad limits. The cube must + remain within 6 cm of TCP. Arm endpoint/tracking tolerances are 0.05 / 0.08 rad. + Motion limits and path/duration quality checks remain mandatory. + +These are conservative sampled geometric checks plus actual discrete-physics +contact evidence, not continuous collision detection. The integration currently +supports one unscaled fixed-base URDF robot and cuboid rigid scene objects. +It requires the optional ``python-fcl``, ``trimesh`` and ``yourdfpy`` dependencies. +Contact-aware Gym execution requires equivalent physics-substep evidence and +is rejected at construction today. + +Between rounds, the entire robot/cube state is restored and verified, including +current joints, passive coordinates, velocities, drive targets and efforts. +Every returned episode has T actual controller targets and T+1 measured +observations/timestamps. ``commanded_joint_indices`` identifies actuated columns +in the full ordered joint vector; passive target columns are retained as read +from the controller, while collision checks use actual passive positions. + +The collection saves ``generation_report.json`` (attempts, rejections, validation +and commit audit), ``pickup_report.json`` (contact profile, prepared rounds and +video frames), and ``manifest.json`` (only read-back-confirmed shards). Each shard +contains LeRobot training frames plus ``episode.json`` and ``terminal.npz``. +Object/TCP 4x4 matrices are flattened in C order into 16-element LeRobot numeric +features; ``observation_shapes`` records their original layout. Terminal matrices +retain their 4x4 shape. ``preview.mp4`` includes accepted and rejected attempts; +it is a four-panel record of actual observations, and each round restarts its +trail and time display. Video is separate from the numeric training dataset. + +``source.kind=atomic`` here identifies an offline compilation source. Exporting +its trajectory does not commit projected symbolic effects or execute the atomic +recovery/tracking state machine. A full runtime source adapter, the four-way +handwritten/atomic × sim/Gym qualification matrix, general via-point planning, +and YAML-based deployment remain subsequent implementation milestones. diff --git a/docs/source/quick_start/install.md b/docs/source/quick_start/install.md index ae6e17c36..172fba5a3 100644 --- a/docs/source/quick_start/install.md +++ b/docs/source/quick_start/install.md @@ -238,7 +238,7 @@ EmbodiChain installation: Linux, Python 3.10--3.13, a supported NVIDIA GPU with at least 4 GB VRAM, and a driver that supports CUDA 12 or newer. See [NVIDIA's official installation guide](https://nvlabs.github.io/curobo/latest/getting-started/installation.html) for the current compatibility requirements, and see -[cuRobo V2 Planner](../overview/sim/planners/curobo_planner.md) for EmbodiChain +[cuRobo V2 Planner](../overview/sim/motion/planners/curobo_planner.md) for EmbodiChain configuration and usage. cuRobo planning always runs on CUDA, but the SimulationManager physics device may be either CPU or CUDA. diff --git a/docs/source/resources/robot/franka_panda.md b/docs/source/resources/robot/franka_panda.md index e116b0640..bd219ff1a 100644 --- a/docs/source/resources/robot/franka_panda.md +++ b/docs/source/resources/robot/franka_panda.md @@ -104,5 +104,5 @@ See :doc:`dual_arm` for the available mount presets and configuration options. ## See Also -- :doc:`/overview/sim/solvers/index` - IK solver reference +- :doc:`/overview/sim/motion/solvers/index` - IK solver reference - :doc:`/guides/add_robot` - Adding a new robot (quick reference) diff --git a/docs/source/resources/robot/ur_robot.md b/docs/source/resources/robot/ur_robot.md index 4056dcc3e..8319d61d7 100644 --- a/docs/source/resources/robot/ur_robot.md +++ b/docs/source/resources/robot/ur_robot.md @@ -89,4 +89,4 @@ robot = sim.add_robot(cfg=cfg) ## See Also - :doc:`/guides/add_robot` - Adding a new robot (quick reference) -- :doc:`/overview/sim/solvers/index` - IK solver reference +- :doc:`/overview/sim/motion/solvers/index` - IK solver reference diff --git a/docs/source/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index 02c0aa8fa..8bc870847 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -156,7 +156,7 @@ total trajectory duration four times longer. The script checks that relationship before replaying the fast and slow trajectories in sequence. The ``motion_generator`` variable in the snippets below is a configured -:class:`~embodichain.lab.sim.planners.MotionGenerator`; its robot, planner, +:class:`~embodichain.lab.sim.motion.planners.MotionGenerator`; its robot, planner, device, cache, and collision world become the resources owned by the engine. Control-part commands diff --git a/docs/source/tutorial/motion_gen.rst b/docs/source/tutorial/motion_gen.rst index 9d550c454..c21a388b3 100644 --- a/docs/source/tutorial/motion_gen.rst +++ b/docs/source/tutorial/motion_gen.rst @@ -3,7 +3,7 @@ Motion Generator ================ -.. currentmodule:: embodichain.lab.sim.planners.motion_generator +.. currentmodule:: embodichain.lab.sim.motion.planners.motion_generator The ``MotionGenerator`` class in EmbodiChain provides a unified and extensible interface for robot trajectory planning. It supports time-optimal trajectory generation (currently via TOPPRA), joint/Cartesian interpolation, and is designed for easy integration with RL, imitation learning, and classical control scenarios. @@ -34,9 +34,9 @@ Typical Usage .. code-block:: python - from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg - from embodichain.lab.sim.planners.toppra_planner import ToppraPlanOptions - from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod, MoveType + from embodichain.lab.sim.motion.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg + from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanOptions + from embodichain.lab.sim.motion.planners.utils import PlanState, TrajectorySampleMethod, MoveType # Assume you have a robot instance and arm_name # Constraints are now specified in ToppraPlanOptions, not in ToppraPlannerCfg @@ -87,7 +87,7 @@ API Reference .. code-block:: python - from embodichain.lab.sim.planners.toppra_planner import ToppraPlanOptions + from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanOptions motion_cfg = MotionGenCfg( planner_cfg=ToppraPlannerCfg( diff --git a/docs/source/tutorial/solver.rst b/docs/source/tutorial/solver.rst index 10fad00cd..a8a17da97 100644 --- a/docs/source/tutorial/solver.rst +++ b/docs/source/tutorial/solver.rst @@ -4,7 +4,7 @@ Create a solver =============== -.. currentmodule:: embodichain.lab.sim.solvers +.. currentmodule:: embodichain.lab.sim.motion.solvers Overview ~~~~~~~~ diff --git a/embodichain/cli/workspace_cache.py b/embodichain/cli/workspace_cache.py index 7595c7be4..e6a20bfe8 100644 --- a/embodichain/cli/workspace_cache.py +++ b/embodichain/cli/workspace_cache.py @@ -59,7 +59,7 @@ def main(argv: Sequence[str] | None = None) -> None: ) args = parser.parse_args(argv) - from embodichain.lab.sim.workspace.caches.cache_utils import ( + from embodichain.lab.sim.motion.workspace.caches.cache_utils import ( clean_all_sessions, clean_session, list_sessions, diff --git a/embodichain/lab/gym/envs/action_bank/configurable_action.py b/embodichain/lab/gym/envs/action_bank/configurable_action.py index 964c2b1c6..01b37c7bd 100644 --- a/embodichain/lab/gym/envs/action_bank/configurable_action.py +++ b/embodichain/lab/gym/envs/action_bank/configurable_action.py @@ -1267,7 +1267,7 @@ def plan_trajectory( vis: bool = False, **kwargs, ) -> np.ndarray: - from embodichain.lab.sim.planners import ( + from embodichain.lab.sim.motion.planners import ( MoveType, PlanState, MotionGenerator, diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index 536a50521..8df2ed3a3 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -17,7 +17,8 @@ from __future__ import annotations import math -from collections.abc import Mapping +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager from numbers import Integral, Real import torch @@ -136,6 +137,11 @@ def __init__( **kwargs, ): self.cfg = cfg + self._generation_lease_owner: object | None = None + self._generation_epoch = 0 + self._generation_prepared = False + self._generation_preparing = False + self._generation_no_auto_reset = False # the number of envs to be simulated in parallel. self._num_envs = self.cfg.num_envs @@ -820,6 +826,108 @@ def _step_action(self, action: EnvAction) -> EnvAction: """ pass + @property + def generation_epoch(self) -> int: + """Return the revision used to invalidate generation runtime bindings. + + Acquisition, release, each controlled preparation attempt, and normal + resets advance this revision, including preparations that fail. + """ + return getattr(self, "_generation_epoch", 0) + + def acquire_generation_lease(self, owner: object) -> None: + """Reserve this complete environment batch for controlled generation. + + The environment and simulator share one owner identity. The same owner + may acquire its lease repeatedly. While reserved, normal resets are + rejected and automatic resets are disabled. The caller must prepare a + generation episode before stepping and serialize host access. + + Args: + owner: Identity token retained until release; must not be ``None``. + + Raises: + ValueError: If the owner is ``None``. + RuntimeError: If another owner already holds the lease. + """ + if owner is None: + raise ValueError("A generation lease owner must not be None.") + current = getattr(self, "_generation_lease_owner", None) + sim_owner = getattr(self.sim, "_trajectory_generation_owner", None) + if current is owner: + self._require_generation_lease(owner) + return + if current is not None or sim_owner is not None: + raise RuntimeError( + "The environment or simulator already has a generation lease owner." + ) + self.sim._trajectory_generation_owner = owner + self._generation_lease_owner = owner + self._generation_epoch = self.generation_epoch + 1 + self._generation_prepared = False + self._generation_no_auto_reset = True + + def _require_generation_lease(self, owner: object) -> None: + """Require the same lease identity on the environment and simulator.""" + if ( + owner is None + or getattr(self, "_generation_lease_owner", None) is not owner + or getattr(self.sim, "_trajectory_generation_owner", None) is not owner + ): + raise RuntimeError("The caller does not own the generation lease.") + + def release_generation_lease(self, owner: object) -> None: + """Release generation ownership without saving or resetting an episode. + + Args: + owner: The identity token used to acquire the lease. + + Raises: + RuntimeError: If the caller is not the owner or preparation is active. + """ + self._require_generation_lease(owner) + if getattr(self, "_generation_preparing", False): + raise RuntimeError("Cannot release a generation lease during preparation.") + if getattr(self, "_generation_command_observer", None) is not None: + raise RuntimeError( + "Cannot release a generation lease during command observation." + ) + self.sim._trajectory_generation_owner = None + self._generation_lease_owner = None + self._generation_epoch = self.generation_epoch + 1 + self._generation_prepared = False + self._generation_no_auto_reset = False + + @contextmanager + def observe_generation_commands( + self, owner: object, observer: Callable[[], None] + ) -> Iterator[None]: + """Observe successfully submitted controller commands before physics. + + The callback runs after ``_step_action`` returns and before simulation + advances. It can read actual robot targets without action postprocessing + changing the evidence. Exceptions propagate after command submission; + callers must count the attempt and safely stop the robot. Only the + generation lease owner may install one serialized observer. + + Args: + owner: Current generation lease identity. + observer: Callback invoked once per submitted batched command. + + Yields: + Control while the observer is installed. + """ + self._require_generation_lease(owner) + if not callable(observer): + raise TypeError("The generation command observer must be callable.") + if getattr(self, "_generation_command_observer", None) is not None: + raise RuntimeError("A generation command observer is already installed.") + self._generation_command_observer = (owner, observer) + try: + yield + finally: + self._generation_command_observer = None + def reset( self, seed: int | None = None, options: dict | None = None ) -> Tuple[EnvObs, Dict]: @@ -831,7 +939,15 @@ def reset( Returns: A tuple containing the observations and infos. + + Raises: + RuntimeError: If a generation lease owns this environment batch. """ + if getattr(self, "_generation_lease_owner", None) is not None: + raise RuntimeError( + "Normal reset is disabled while a generation lease is held." + ) + self._generation_epoch = self.generation_epoch + 1 if seed is not None: seed = self._set_seed(seed) super().reset(seed=seed) @@ -895,13 +1011,27 @@ def step( Returns: A tuple contraining the observation, reward, terminated, truncated, and info dictionary. + + Raises: + RuntimeError: If generation owns the batch without a successfully + prepared episode. """ + if getattr(self, "_generation_lease_owner", None) is not None and not getattr( + self, "_generation_prepared", False + ): + raise RuntimeError("Prepare a valid generation episode before stepping.") + with self._profiler.section("step", is_root=True): with self._profiler.section("preprocess_action"): action = self._preprocess_action(action=action) with self._profiler.section("step_action"): action = self._step_action(action=action) + command_observer = getattr(self, "_generation_command_observer", None) + if command_observer is not None: + observer_owner, observer = command_observer + self._require_generation_lease(observer_owner) + observer() with self._profiler.section("sim_update"): self.sim.update(self.physics_dt, self.cfg.sim_steps_per_control) @@ -957,6 +1087,7 @@ def step( if not ( getattr(self, "_replay_no_auto_reset", False) or getattr(self, "_demo_no_auto_reset", False) + or getattr(self, "_generation_no_auto_reset", False) ): reset_env_ids = dones.nonzero(as_tuple=False).squeeze(-1) if len(reset_env_ids) > 0: diff --git a/embodichain/lab/gym/envs/demo.py b/embodichain/lab/gym/envs/demo.py index 6414bf8ea..708a2ea01 100644 --- a/embodichain/lab/gym/envs/demo.py +++ b/embodichain/lab/gym/envs/demo.py @@ -649,18 +649,36 @@ def resolve_demo_segments(env: Any, **kwargs: Any) -> Iterable[DemoSegment]: else (DemoSegment(actions, name="legacy", metadata={"segment_count": 1}),) ) + return _validated_demo_segments( + segments, + fallback_instruction=_dataset_instruction(env), + source="create_demo_segments()", + ) + + +def _validated_demo_segments( + segments: Iterable[DemoSegment] | DemoSegment | None, + *, + fallback_instruction: str, + source: str, +) -> Iterable[DemoSegment]: + """Normalize a segment source without eagerly consuming lazy plans.""" if segments is None: return () if isinstance(segments, DemoSegment): segments = (segments,) - - fallback_instruction = _dataset_instruction(env) + try: + iterator = iter(segments) + except TypeError as exc: + raise TypeError( + f"{source} must be a DemoSegment or an iterable of DemoSegment objects." + ) from exc def _validate() -> Iterable[DemoSegment]: - for segment in segments: + for segment in iterator: if not isinstance(segment, DemoSegment): raise TypeError( - "create_demo_segments() must yield DemoSegment objects, " + f"{source} must yield DemoSegment objects, " f"got {type(segment).__name__}." ) if segment.instruction is None: @@ -673,14 +691,17 @@ def _validate() -> Iterable[DemoSegment]: def execute_demo_episode( env: Any, *, + segments: Iterable[DemoSegment] | DemoSegment | None = None, episode_index: int = 0, execution_cfg: DemoExecutionCfg | None = None, attempt_id: int = 0, should_stop: StopPredicate | None = None, progress: ProgressWrapper | None = None, + step_observer: Callable[[Any, tuple[bool, ...]], None] | None = None, + row_step_limits: tuple[int, ...] | None = None, **plan_kwargs: Any, ) -> DemoEpisodeResult: - """Plan and execute every segment in one environment episode. + """Execute supplied segments or plan one environment demonstration episode. Auto-reset is suspended for the duration of execution. The caller owns the transaction boundary and must explicitly call ``env.reset()`` to commit a @@ -689,18 +710,36 @@ def execute_demo_episode( Args: env: Gym environment or wrapper. + segments: Explicit candidate segments, consumed lazily without calling + the environment's planning methods. ``None`` retains task-owned + planning. Explicit segments cannot be combined with planning + arguments. episode_index: Logical episode identifier used in metadata and logs. execution_cfg: Collector-owned output settings. Defaults to continuous episode persistence. attempt_id: Zero-based identifier for this collection attempt. should_stop: Optional callback checked before every action. progress: Optional wrapper such as ``tqdm`` for action iterables. + step_observer: Optional callback receiving the unchanged ``env.step`` + result and the rows active for that transition. It runs before + terminal handling, without querying observations again. + row_step_limits: Optional full-batch action counts for supplied ragged + trajectories. Zero skips a row. Finished rows stop recording and + receive normal demo hold commands while others finish; their task + success is checked at the final batch boundary. ``None`` preserves + ordinary segment execution. **plan_kwargs: Arguments forwarded to the task's planning method. Returns: A :class:`DemoEpisodeResult` describing segment spans and terminal state. + + Raises: + TypeError: If a segment source contains values other than DemoSegment. + ValueError: If explicit segments and planning arguments are combined. """ + if segments is not None and plan_kwargs: + raise ValueError("Explicit segments cannot be combined with plan_kwargs.") if execution_cfg is None: execution_cfg = DemoExecutionCfg() elif not isinstance(execution_cfg, DemoExecutionCfg): @@ -710,6 +749,18 @@ def execute_demo_episode( target = _env_target(env) num_envs = int(getattr(target, "num_envs", 1)) + if step_observer is not None and not callable(step_observer): + raise TypeError("step_observer must be callable or None.") + if row_step_limits is not None: + row_step_limits = tuple(row_step_limits) + if ( + segments is None + or len(row_step_limits) != num_envs + or any(type(value) is not int or value < 0 for value in row_step_limits) + ): + raise ValueError( + "row_step_limits require explicit segments and one nonnegative action count per row." + ) begin_episode = _get_env_callable(env, "_begin_demo_episode_recording") begin_segment = _get_env_callable(env, "_begin_demo_segment_recording") end_segment = _get_env_callable(env, "_end_demo_segment_recording") @@ -719,7 +770,10 @@ def execute_demo_episode( set_active_mask = _get_env_callable(env, "_set_demo_active_mask") success_fn = _get_env_callable(env, "is_task_success") - active = [True] * num_envs + active = [ + row_step_limits is None or row_step_limits[row] > 0 for row in range(num_envs) + ] + row_exhausted = [False] * num_envs def publish_active_mask() -> None: """Publish executor liveness to recording hooks and action masking.""" @@ -757,7 +811,15 @@ def publish_active_mask() -> None: try: segment_count = 0 - segments = iter(resolve_demo_segments(env, **plan_kwargs)) + segment_iterator = iter( + resolve_demo_segments(env, **plan_kwargs) + if segments is None + else _validated_demo_segments( + segments, + fallback_instruction=_dataset_instruction(env), + source="segments", + ) + ) while any(active): if should_stop is not None and should_stop(): fatal_reason = "interrupted" @@ -768,7 +830,7 @@ def publish_active_mask() -> None: publish_active_mask() break try: - segment = next(segments) + segment = next(segment_iterator) except StopIteration: break @@ -847,7 +909,8 @@ def publish_active_mask() -> None: active_before_step = tuple(active) try: - _, _, terminated_value, truncated_value, info = env.step(action) + step_result = env.step(action) + _, _, terminated_value, truncated_value, info = step_result except Exception as exc: action_error = exc actions_exhausted = False @@ -859,6 +922,14 @@ def publish_active_mask() -> None: for env_id, was_active in enumerate(active_before_step): if was_active: lengths[env_id] += 1 + if step_observer is not None: + try: + step_observer(step_result, active_before_step) + except Exception as exc: + action_error = exc + actions_exhausted = False + segment_reason = "step_observation_failed" + break step_terminated = _as_bool_tuple(terminated_value, num_envs) step_truncated = _as_bool_tuple(truncated_value, num_envs) @@ -926,11 +997,17 @@ def publish_active_mask() -> None: actions_exhausted = False break + if row_step_limits is not None: + for env_id, is_active in enumerate(active): + if is_active and lengths[env_id] >= row_step_limits[env_id]: + active[env_id] = False + row_exhausted[env_id] = True + terminal_reasons[env_id] = "trajectory_exhausted" publish_active_mask() if not any(active): # Every row reached episode-level success. Stop this segment # and do not request another lazy segment. - actions_exhausted = False + actions_exhausted = any(row_exhausted) break if should_stop is not None and should_stop(): actions_exhausted = False @@ -1037,10 +1114,13 @@ def publish_active_mask() -> None: continue if completed_by_env[env_id] and success[env_id]: segment_successes[env_id] = True - elif active[env_id] and validation[env_id]: + elif (active[env_id] or row_exhausted[env_id]) and validation[ + env_id + ]: segment_successes[env_id] = True - elif active[env_id]: + elif active[env_id] or row_exhausted[env_id]: validation_failed = True + row_exhausted[env_id] = False segment_failure_reasons[env_id] = "segment_validation_failed" terminal_reasons[env_id] = "segment_validation_failed" @@ -1115,7 +1195,7 @@ def publish_active_mask() -> None: terminal_reasons[env_id] = fatal_reason active[env_id] = False publish_active_mask() - elif fatal_reason is None and any(active): + elif fatal_reason is None and (any(active) or any(row_exhausted)): # Normal plan exhaustion validates only rows that have not already # reached sticky episode success. Legacy expert tasks use # is_task_success() for this final validation. @@ -1126,7 +1206,7 @@ def publish_active_mask() -> None: ) final_success = _as_bool_tuple(success_source, num_envs) for env_id, is_active in enumerate(active): - if not is_active: + if not is_active and not row_exhausted[env_id]: continue success[env_id] = final_success[env_id] completed_by_env[env_id] = final_success[env_id] diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index b8de4176e..1ca67b4de 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -16,7 +16,7 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Callable, Mapping from math import log from functools import wraps from datetime import datetime @@ -83,6 +83,7 @@ from embodichain.data.constants import EMBODICHAIN_DEFAULT_DATA_ROOT if TYPE_CHECKING: + from embodichain.lab.sim.motion.trajectory_augmentation import ValidationResult from embodichain.lab.task_program import CompiledTaskProgram, TaskProgramCfg from embodichain.lab.task_program.integrations import ( TaskProgramAdapterFactory, @@ -574,6 +575,85 @@ def reset( self._seed_recording_state(obs, reset_ids) return obs, info + def prepare_generation_episode( + self, + owner: object, + *, + prepare: Callable[[], None], + restore: Callable[[], None], + settle: Callable[[], None], + verify: Callable[[], ValidationResult], + ) -> tuple[EnvObs, Dict[str, Any]]: + """Discard buffered evidence and prepare the whole leased batch. + + Freeze any previous episode before calling this method. Host callbacks + own deterministic task/controller preparation, physical restoration, + settling, and initial-state verification. They must not call ``step`` + or ``reset``. Reset/interval events and environment seeding are never + run here. Standard managers and observation histories are reset after + settling, before verification checks their initial state. Recording is + seeded only after successful verification and fresh observation capture. + + Every attempt invalidates earlier generation epochs. A failed attempt + leaves stepping disabled until another preparation succeeds. + + Args: + owner: The current generation lease identity. + prepare: Initialize deterministic task and controller state. + restore: Restore the host's complete captured initial state. + settle: Advance the host through its explicit settling policy. + verify: Return nonempty, all-passed initial-state evidence. + + Returns: + The prepared observation and task information for the entire batch. + + Raises: + RuntimeError: If the lease is invalid, preparation is active, or + initial-state verification does not pass. + TypeError: If callbacks are not callable or verification returns + something other than ``ValidationResult``. + """ + from embodichain.lab.sim.motion.trajectory_augmentation import ValidationResult + + self._require_generation_lease(owner) + if getattr(self, "_generation_preparing", False): + raise RuntimeError("Generation episode preparation is already active.") + self._generation_epoch = self.generation_epoch + 1 + self._generation_prepared = False + self._active_task_program_bridge = None + if not all( + callable(callback) for callback in (prepare, restore, settle, verify) + ): + raise TypeError("Generation preparation requires four callable callbacks.") + env_ids = torch.arange(self.num_envs, dtype=torch.long, device=self.device) + self._generation_preparing = True + try: + self._finish_camera_recordings(env_ids, save_data=False) + self._clear_episode_recording_state(env_ids) + self._elapsed_steps.zero_() + self._task_success.zero_() + self._traj_raw_action = None + prepare() + restore() + settle() + self._reset_episode_managers(env_ids) + validation = verify() + if not isinstance(validation, ValidationResult): + raise TypeError( + "Initial-state verification must return ValidationResult." + ) + if not validation.accepted: + raise RuntimeError( + "Generation initial-state verification did not pass." + ) + obs = self.get_obs() + info = self.get_info() + self._seed_recording_state(obs, env_ids) + self._generation_prepared = True + return obs, info + finally: + self._generation_preparing = False + def _seed_recording_state(self, obs: EnvObs, env_ids: torch.Tensor) -> None: """Seed all enabled recorders from the current environment state.""" self._seed_expert_observations(obs, env_ids) @@ -943,22 +1023,9 @@ def _initialize_episode( env_ids=env_ids_to_save, ) - # Save recorded camera data before resetting - if self.cfg.events and self.event_manager is not None: - from embodichain.lab.gym.envs.managers.record import record_camera_data - - with self._profiler.section("record_camera_save"): - for mode_cfgs in self.event_manager._mode_functor_cfgs.values(): - for functor_cfg in mode_cfgs: - if isinstance(functor_cfg.func, record_camera_data): - if save_data: - functor_cfg.func.save_and_clear( - env_ids=env_ids_to_process - ) - else: - functor_cfg.func.discard_and_clear( - env_ids=env_ids_to_process - ) + EmbodiedEnv._finish_camera_recordings( + self, env_ids_to_process, save_data=save_data + ) # Auto-save + reset the per-env trajectory buffer for environments being # reset. Use getattr so this no-ops on envs/subclasses that don't allocate @@ -973,26 +1040,54 @@ def _initialize_episode( for env_id in env_ids_to_process.tolist(): self._save_trajectory_for_env(env_id) + EmbodiedEnv._clear_episode_recording_state(self, env_ids_to_process) + + # apply events such as randomization for environments that need a reset + if self.cfg.events: + if "reset" in self.event_manager.available_modes: + with self._profiler.section("event_reset"): + self.event_manager.apply(mode="reset", env_ids=env_ids) + + EmbodiedEnv._reset_episode_managers(self, env_ids) + + def _finish_camera_recordings( + self, env_ids: torch.Tensor, *, save_data: bool + ) -> None: + """Finish selected camera buffers without changing event streams.""" + if self.cfg.events and self.event_manager is not None: + from embodichain.lab.gym.envs.managers.record import record_camera_data + + with self._profiler.section("record_camera_save"): + for mode_cfgs in self.event_manager._mode_functor_cfgs.values(): + for functor_cfg in mode_cfgs: + if isinstance(functor_cfg.func, record_camera_data): + if save_data: + functor_cfg.func.save_and_clear(env_ids=env_ids) + else: + functor_cfg.func.discard_and_clear(env_ids=env_ids) + + def _clear_episode_recording_state(self, env_ids: torch.Tensor) -> None: + """Discard episode buffers and annotations without saving or events.""" _traj_steps = getattr(self, "_traj_steps", None) if _traj_steps is not None: - _traj_steps[env_ids_to_process] = 0 + _traj_steps[env_ids] = 0 # Clear episode buffers only after every recorder has consumed them. if self.rollout_buffer is not None and self._rollout_buffer_mode != "rl": - self._clear_expert_rollout_rows(env_ids_to_process) + self._clear_expert_rollout_rows(env_ids) rollout_steps = getattr(self, "rollout_steps", None) if rollout_steps is not None: - rollout_ids = env_ids_to_process.to(rollout_steps.device) + rollout_ids = env_ids.to(rollout_steps.device) rollout_steps[rollout_ids] = 0 self.current_rollout_step = int(rollout_steps.max().item()) episode_metadata = getattr(self, "_demo_episode_metadata", None) if episode_metadata is not None: - for env_id in env_ids_to_process.cpu().tolist(): + for env_id in env_ids.cpu().tolist(): episode_metadata[env_id] = self._new_demo_episode_metadata(env_id) active_segment_ids = getattr(self, "_demo_active_segment_ids", None) if active_segment_ids is not None: - demo_ids = env_ids_to_process.to(active_segment_ids.device) + demo_ids = env_ids.to(active_segment_ids.device) active_segment_ids[demo_ids] = 0 self._demo_active_mask[demo_ids] = True self._demo_segment_participants[demo_ids] = False @@ -1000,14 +1095,12 @@ def _initialize_episode( self._demo_active_rollout_start_steps[demo_ids] = 0 self._demo_steps[demo_ids] = 0 - self.episode_success_status[env_ids_to_process] = False - - # apply events such as randomization for environments that need a reset - if self.cfg.events: - if "reset" in self.event_manager.available_modes: - with self._profiler.section("event_reset"): - self.event_manager.apply(mode="reset", env_ids=env_ids) + self.episode_success_status[env_ids] = False + def _reset_episode_managers( + self, env_ids: Sequence[int] | torch.Tensor | None + ) -> None: + """Reset observation history, reward, and dataset functors only.""" # reset observation manager for environments that need a reset # This clears any cached data in observation functors (e.g., physics attributes) if self.cfg.observations: diff --git a/embodichain/lab/scripts/analyze_workspace.py b/embodichain/lab/scripts/analyze_workspace.py index 7f1649c14..bc28c1b72 100644 --- a/embodichain/lab/scripts/analyze_workspace.py +++ b/embodichain/lab/scripts/analyze_workspace.py @@ -71,7 +71,7 @@ if TYPE_CHECKING: from embodichain.lab.sim.cfg import RobotCfg - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( WorkspaceAnalyzerConfig, ) @@ -298,7 +298,7 @@ def _build_asset_robot_cfg( given without ``--urdf``. """ from embodichain.lab.sim.cfg import RobotCfg - from embodichain.lab.sim.solvers import ( + from embodichain.lab.sim.motion.solvers import ( PinkSolverCfg, PinocchioSolverCfg, PytorchSolverCfg, @@ -365,7 +365,7 @@ def build_analyzer_config( """ import torch - from embodichain.lab.sim.workspace.configs import ( + from embodichain.lab.sim.motion.workspace.configs import ( CacheConfig, DimensionConstraint, SamplingConfig, @@ -373,7 +373,7 @@ def build_analyzer_config( VisualizationConfig, VisualizationType, ) - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( AnalysisMode, WorkspaceAnalyzerConfig, ) @@ -592,7 +592,7 @@ def preview_cache(args: argparse.Namespace, analyzer: WorkspaceAnalyzer) -> None ``args.hide_unreachable``. analyzer: Analyzer attached to the loaded robot and simulation. """ - from embodichain.lab.sim.workspace.caches.results_cache import ( + from embodichain.lab.sim.motion.workspace.caches.results_cache import ( DEFAULT_RESULTS_CACHE_DIR, deserialize_results, ) @@ -638,7 +638,7 @@ def main(args: argparse.Namespace) -> None: import torch from embodichain.lab.sim.sim_manager import SimulationManager - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( WorkspaceAnalyzer, ) @@ -1069,7 +1069,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: # Resolve cache dir default here (after parse) to keep --help output clean. if args.cache_dir is None and not args.no_cache: - from embodichain.lab.sim.workspace.caches import ( + from embodichain.lab.sim.motion.workspace.caches import ( DEFAULT_RESULTS_CACHE_DIR, ) diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index 69baaa9c2..2684477c5 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -19,7 +19,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Mapping +from collections.abc import Callable, Mapping from copy import deepcopy from dataclasses import dataclass, field, replace from functools import cached_property @@ -64,7 +64,7 @@ if TYPE_CHECKING: from embodichain.lab.sim.objects import Robot - from embodichain.lab.sim.planners import MotionGenerator + from embodichain.lab.sim.motion.planners import MotionGenerator from .runtime import ActionPlanningServices from .state import PlanningContext @@ -406,11 +406,22 @@ def plan( Returns: Scene-bound action plan with expected, uncommitted effects. """ + return self._plan_with_provider(request, context, self._plan) + + def _plan_with_provider( + self, + request: ResolvedActionRequest[GoalT, OptionsT], + context: PlanningContext, + plan_provider: Callable[ + [ResolvedActionRequest[GoalT, OptionsT], PlanningContext], ActionPlan + ], + ) -> ActionPlan: + """Apply framework validation to an ordinary or supplied initial plan.""" self.require_goal(request) prepared = self._prepare_request(request, context) - plan = self._plan(prepared, context) + plan = plan_provider(prepared, context) if not isinstance(plan, ActionPlan): - raise TypeError("AtomicAction._plan() must return an ActionPlan.") + raise TypeError("Atomic action plan provider must return an ActionPlan.") return replace( plan, commands=self._authorize_command_targets(prepared, plan.commands), diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index d86539e79..d15340634 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -19,7 +19,7 @@ from __future__ import annotations from types import MappingProxyType -from typing import Iterable, Mapping, TYPE_CHECKING +from typing import Callable, Iterable, Mapping, TYPE_CHECKING import torch @@ -36,7 +36,7 @@ if TYPE_CHECKING: from embodichain.lab.sim.objects import Robot - from embodichain.lab.sim.planners import MotionGenerator + from embodichain.lab.sim.motion.planners import MotionGenerator from embodichain.toolkits.graspkit import GraspPoseGenerator from .execution import ExecutionSession @@ -349,6 +349,10 @@ def _plan_request( self, request: ResolvedActionRequest, context: PlanningContext | None = None, + *, + plan_provider: ( + Callable[[ResolvedActionRequest, PlanningContext], ActionPlan] | None + ) = None, ) -> ActionPlan: """Plan an already-resolved request without rebuilding its snapshot. @@ -359,6 +363,8 @@ def _plan_request( Args: request: Immutable request previously returned by :meth:`_resolve`. context: Optional latest planning state; captured when omitted. + plan_provider: Initial-plan materializer supplied by a new session. + Recovery leaves this unset and invokes the registered planner. Returns: Validated side-effect-free action plan. @@ -372,7 +378,11 @@ def _plan_request( ) current = self.initial_context() if context is None else context self._validate_context(current) - plan = action.plan(request, current) + plan = ( + action.plan(request, current) + if plan_provider is None + else action._plan_with_provider(request, current, plan_provider) + ) self._validate_plan(plan, current, request) return plan @@ -524,6 +534,9 @@ def start( context: PlanningContext | None = None, *, eligible_mask: torch.Tensor | None = None, + initial_plan_provider: ( + Callable[[ResolvedActionRequest, PlanningContext], ActionPlan] | None + ) = None, ) -> ExecutionSession: """Start incremental execution for a grounded invocation sequence. @@ -534,6 +547,13 @@ def start( eligible_mask: Optional per-environment cohort allowed to execute. Ineligible rows remain excluded for the whole session. All rows are eligible when omitted. + initial_plan_provider: Optional materializer for the first invocation's + initial plan. It receives a newly resolved request with current + collision options and the initial measured context. Its result + passes ordinary plan, endpoint, tracking, and phase-gate checks. + Later invocations and recovery use the registered skill planner. + Rebuild candidate bindings and effects from these inputs; do not + reuse a plan or execution session across environment resets. Returns: Stateful execution session advanced by ``session.tick(...)``. @@ -546,6 +566,7 @@ def start( tuple(invocations), initial, eligible_mask=eligible_mask, + initial_plan_provider=initial_plan_provider, ) def _validate_context(self, context: PlanningContext) -> None: diff --git a/embodichain/lab/sim/atomic_actions/execution.py b/embodichain/lab/sim/atomic_actions/execution.py index 0af767e1a..abdd50cfb 100644 --- a/embodichain/lab/sim/atomic_actions/execution.py +++ b/embodichain/lab/sim/atomic_actions/execution.py @@ -21,7 +21,7 @@ from dataclasses import dataclass, replace from enum import Enum import math -from typing import TYPE_CHECKING +from typing import Callable, TYPE_CHECKING import torch @@ -325,9 +325,14 @@ def __init__( context: PlanningContext, *, eligible_mask: torch.Tensor | None = None, + initial_plan_provider: ( + Callable[[ResolvedActionRequest, PlanningContext], ActionPlan] | None + ) = None, ) -> None: if not invocations: raise ValueError("ExecutionSession requires at least one invocation.") + if initial_plan_provider is not None and not callable(initial_plan_provider): + raise TypeError("initial_plan_provider must be callable.") engine._validate_context(context) self._engine = engine self._requests: tuple[ResolvedActionRequest, ...] = tuple( @@ -388,7 +393,11 @@ def __init__( ) self._queued_events: list[ExecutionEvent] = [] if self._status is ExecutionStatus.RUNNING: - self._plan_current(context, ExecutionEventKind.ACTION_PLANNED) + self._plan_current( + context, + ExecutionEventKind.ACTION_PLANNED, + plan_provider=initial_plan_provider, + ) else: self._queued_events.append( self._event( @@ -1145,10 +1154,14 @@ def _plan_current( self, context: PlanningContext, event_kind: ExecutionEventKind, + *, + plan_provider: ( + Callable[[ResolvedActionRequest, PlanningContext], ActionPlan] | None + ) = None, ) -> None: """Plan the current invocation from the latest observation.""" request = self._requests[self._invocation_index] - plan = self._engine._plan_request(request, context) + plan = self._engine._plan_request(request, context, plan_provider=plan_provider) self._install_plan(plan, context, event_kind) def _install_plan( diff --git a/embodichain/lab/sim/atomic_actions/plans.py b/embodichain/lab/sim/atomic_actions/plans.py index 9ec059a3f..d2d40228e 100644 --- a/embodichain/lab/sim/atomic_actions/plans.py +++ b/embodichain/lab/sim/atomic_actions/plans.py @@ -26,7 +26,7 @@ import torch -from embodichain.lab.sim.planners.utils import normalize_success_mask +from embodichain.lab.sim.motion.planners.utils import normalize_success_mask from .effects import StateDelta from .policies import RecoveryPolicy diff --git a/embodichain/lab/sim/atomic_actions/policies.py b/embodichain/lab/sim/atomic_actions/policies.py index 668f1669e..0599aff67 100644 --- a/embodichain/lab/sim/atomic_actions/policies.py +++ b/embodichain/lab/sim/atomic_actions/policies.py @@ -26,7 +26,7 @@ if TYPE_CHECKING: import torch - from embodichain.lab.sim.planners import MotionGenOptions, PlanOptions + from embodichain.lab.sim.motion.planners import MotionGenOptions, PlanOptions class DynamicCollisionMode(str, Enum): @@ -116,7 +116,9 @@ def to_motion_gen_options( Returns: Independently owned options for :class:`MotionGenerator`. """ - from embodichain.lab.sim.planners.motion_generator import MotionGenOptions + from embodichain.lab.sim.motion.planners.motion_generator import ( + MotionGenOptions, + ) return MotionGenOptions( strategy=self.strategy, diff --git a/embodichain/lab/sim/atomic_actions/primitives/_helpers.py b/embodichain/lab/sim/atomic_actions/primitives/_helpers.py index b629a79a3..d0be8617c 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/_helpers.py +++ b/embodichain/lab/sim/atomic_actions/primitives/_helpers.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: from embodichain.lab.sim.objects import Robot - from embodichain.lab.sim.planners import MotionGenerator + from embodichain.lab.sim.motion.planners import MotionGenerator from ..policies import MotionPolicy diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index 646d35a7b..7ac8fbdf6 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -574,8 +574,19 @@ def _resolve_grasp_pose( ) num_envs = object_pose.shape[0] n_max_pose = max(r[0].shape[0] for r in grasp_poses_result) - grasp_xpos_padding = torch.zeros( - (num_envs, n_max_pose, 4, 4), dtype=torch.float32, device=self.device + safe_pose = self.robot.compute_fk( + qpos=start_qpos, + name=manipulator.control_part, + to_matrix=True, + ) + if n_max_pose == 0: + return ( + torch.zeros(num_envs, dtype=torch.bool, device=self.device), + safe_pose, + ) + grasp_xpos_padding = safe_pose[:, None].expand(-1, n_max_pose, -1, -1).clone() + candidate_mask = torch.zeros( + (num_envs, n_max_pose), dtype=torch.bool, device=self.device ) grasp_cost_padding = torch.full( (num_envs, n_max_pose), @@ -585,16 +596,47 @@ def _resolve_grasp_pose( ) for i in range(num_envs): n_pose = grasp_poses_result[i][0].shape[0] + if n_pose == 0: + continue grasp_poses = grasp_poses_result[i][0].to( - device=self.device, dtype=torch.float32 + device=self.device, dtype=grasp_xpos_padding.dtype ) grasp_costs = grasp_poses_result[i][1].to( device=self.device, dtype=torch.float32 ) - grasp_xpos_padding[i, :n_pose] = grasp_poses - grasp_cost_padding[i, :n_pose] = grasp_costs - grasp_xpos_padding[i, n_pose:] = grasp_poses[0] - grasp_cost_padding[i, n_pose:] = grasp_costs[0] + valid = torch.isfinite(grasp_costs) & torch.isfinite(grasp_poses).all( + dim=(-2, -1) + ) + # Validate finite transforms before any rotation math or batch IK. + # Invalid candidates remain masked even if the safe placeholder is + # itself reachable by the robot. + checked_poses = torch.where( + valid[:, None, None], grasp_poses, safe_pose[i] + ).to(dtype=torch.float64) + rotation = checked_poses[:, :3, :3] + valid &= torch.isclose( + checked_poses[:, 3], + checked_poses.new_tensor((0.0, 0.0, 0.0, 1.0)), + atol=1.0e-6, + rtol=0.0, + ).all(dim=-1) + valid &= torch.isclose( + rotation.transpose(-2, -1) @ rotation, + torch.eye(3, dtype=rotation.dtype, device=rotation.device), + atol=1.0e-6, + rtol=0.0, + ).all(dim=(-2, -1)) + valid &= torch.isclose( + torch.linalg.det(rotation), + rotation.new_tensor(1.0), + atol=1.0e-6, + rtol=0.0, + ) + grasp_xpos_padding[i, :n_pose] = torch.where( + valid[:, None, None], grasp_poses, safe_pose[i] + ) + grasp_cost_padding[i, :n_pose] = torch.where(valid, grasp_costs, torch.inf) + candidate_mask[i, :n_pose] = valid grasp_xpos_padding, ik_success = self._select_feasible_grasp_variants( grasp_xpos_padding, start_qpos, @@ -603,13 +645,17 @@ def _resolve_grasp_pose( options, approach_direction, ) - grasp_cost_masked = torch.where(ik_success, grasp_cost_padding, 10000.0) + grasp_cost_masked = torch.where( + candidate_mask & ik_success, grasp_cost_padding, torch.inf + ) best_cost, best_idx = grasp_cost_masked.min(dim=1) - is_success = best_cost < 9999.0 + is_success = torch.isfinite(best_cost) best_grasp_xpos = grasp_xpos_padding[ torch.arange(num_envs, device=self.device), best_idx ] - return is_success, best_grasp_xpos + return is_success, torch.where( + is_success[:, None, None], best_grasp_xpos, safe_pose + ) def _select_feasible_grasp_variants( self, @@ -767,8 +813,14 @@ def _compute_batch_candidate_ik( name=manipulator.control_part, joint_seed=flat_seed, ) + success = is_success.to(device=self.device, dtype=torch.bool) + success = success.reshape(num_envs, n_pose * n_variant) + success &= torch.isfinite(qpos).all(dim=-1) + # A failed solver may return NaNs or an arbitrary finite configuration. + # Keep its last valid seed so it cannot contaminate the next stage's IK. + qpos = torch.where(success[..., None], qpos, flat_seed) return ( - is_success.to(device=self.device, dtype=torch.bool).reshape( + success.reshape( num_envs, n_pose, n_variant, diff --git a/embodichain/lab/sim/atomic_actions/runtime.py b/embodichain/lab/sim/atomic_actions/runtime.py index 6ee516bf4..43a251972 100644 --- a/embodichain/lab/sim/atomic_actions/runtime.py +++ b/embodichain/lab/sim/atomic_actions/runtime.py @@ -46,7 +46,7 @@ if TYPE_CHECKING: from embodichain.lab.sim.objects import Robot - from embodichain.lab.sim.planners import MotionGenerator + from embodichain.lab.sim.motion.planners import MotionGenerator class ActionPlanningServices: diff --git a/embodichain/lab/sim/atomic_actions/sim_adapter.py b/embodichain/lab/sim/atomic_actions/sim_adapter.py index 50933d614..1b290f3be 100644 --- a/embodichain/lab/sim/atomic_actions/sim_adapter.py +++ b/embodichain/lab/sim/atomic_actions/sim_adapter.py @@ -46,7 +46,7 @@ if TYPE_CHECKING: from embodichain.lab.sim.objects import RigidObject, Robot - from embodichain.lab.sim.planners import MotionGenerator + from embodichain.lab.sim.motion.planners import MotionGenerator from embodichain.lab.sim.sim_manager import SimulationManager from embodichain.toolkits.graspkit import GraspPoseGenerator diff --git a/embodichain/lab/sim/atomic_actions/trajectory_ops.py b/embodichain/lab/sim/atomic_actions/trajectory_ops.py index 4593ab430..1c666e477 100644 --- a/embodichain/lab/sim/atomic_actions/trajectory_ops.py +++ b/embodichain/lab/sim/atomic_actions/trajectory_ops.py @@ -20,7 +20,7 @@ import torch -from embodichain.lab.sim.planners import MoveType, PlanResult, PlanState +from embodichain.lab.sim.motion.planners import MoveType, PlanResult, PlanState from embodichain.lab.sim.utility.action_utils import interpolate_with_distance from .plans import TimedTrajectory, normalize_success_mask diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index 87549116f..a434af9da 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -47,7 +47,7 @@ from embodichain.utils.utility import key_in_nested_dict from .shapes import ShapeCfg, MeshCfg -from .workspace.cfg import RobotWorkspaceCfg +from .motion.workspace.cfg import RobotWorkspaceCfg # Global default renderer settings for simulation. # @@ -1736,7 +1736,7 @@ def from_dict( @configclass class RobotCfg(ArticulationCfg): - from embodichain.lab.sim.solvers import SolverCfg + from embodichain.lab.sim.motion.solvers import SolverCfg """Configuration for a robot asset in the simulation. """ @@ -1780,7 +1780,7 @@ def from_dict(cls, init_dict: Dict[str, str | float | tuple]) -> RobotCfg: import importlib - solver_module = importlib.import_module("embodichain.lab.sim.solvers") + solver_module = importlib.import_module("embodichain.lab.sim.motion.solvers") cfg = cls() # Create a new instance of the class (cls) for key, value in init_dict.items(): diff --git a/embodichain/lab/sim/motion/__init__.py b/embodichain/lab/sim/motion/__init__.py new file mode 100644 index 000000000..d219539a2 --- /dev/null +++ b/embodichain/lab/sim/motion/__init__.py @@ -0,0 +1,43 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Robot motion solving, planning, workspace analysis, and trajectory augmentation. + +Subpackages load on access so importing solvers during Robot initialization +does not also import planners or workspace analyzers. Importing this package +still follows the normal :mod:`embodichain.lab.sim` initialization lifecycle. +""" + +from __future__ import annotations + +import importlib +from types import ModuleType + +__all__ = ["solvers", "planners", "workspace", "trajectory_augmentation"] + + +def __getattr__(name: str) -> ModuleType: + """Load a public motion subpackage on first access.""" + if name in __all__: + module = importlib.import_module(f"{__name__}.{name}") + globals()[name] = module + return module + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + """Return eagerly and lazily available motion names.""" + return sorted(set(globals()) | set(__all__)) diff --git a/embodichain/lab/sim/planners/__init__.py b/embodichain/lab/sim/motion/planners/__init__.py similarity index 97% rename from embodichain/lab/sim/planners/__init__.py rename to embodichain/lab/sim/motion/planners/__init__.py index d0fd8ef09..45e412e8a 100644 --- a/embodichain/lab/sim/planners/__init__.py +++ b/embodichain/lab/sim/motion/planners/__init__.py @@ -19,6 +19,8 @@ ``BasePlanner`` trajectory planners (TOPPRA, neural, cuRobo) produce joint trajectories from waypoints, and ``MotionGenerator`` wraps a planner with interpolation, IK resolution, and multi-part coordination. """ +from __future__ import annotations + from .utils import * from .base_planner import * from .toppra_planner import * diff --git a/embodichain/lab/sim/planners/base_planner.py b/embodichain/lab/sim/motion/planners/base_planner.py similarity index 100% rename from embodichain/lab/sim/planners/base_planner.py rename to embodichain/lab/sim/motion/planners/base_planner.py diff --git a/embodichain/lab/sim/planners/curobo/curobo_planner.py b/embodichain/lab/sim/motion/planners/curobo/curobo_planner.py similarity index 96% rename from embodichain/lab/sim/planners/curobo/curobo_planner.py rename to embodichain/lab/sim/motion/planners/curobo/curobo_planner.py index d310c26a2..d988472d6 100644 --- a/embodichain/lab/sim/planners/curobo/curobo_planner.py +++ b/embodichain/lab/sim/motion/planners/curobo/curobo_planner.py @@ -48,15 +48,15 @@ from embodichain.utils import configclass, logger from embodichain.utils.math import pose_inv, quat_from_matrix -from embodichain.lab.sim.planners.base_planner import ( +from embodichain.lab.sim.motion.planners.base_planner import ( BasePlanner, BasePlannerCfg, CollisionWorldInfo, PlanOptions, validate_plan_options, ) -from embodichain.lab.sim.planners.curobo.curobo_yaml import _named_rigid_objects -from embodichain.lab.sim.planners.utils import MoveType, PlanResult, PlanState +from embodichain.lab.sim.motion.planners.curobo.curobo_yaml import _named_rigid_objects +from embodichain.lab.sim.motion.planners.utils import MoveType, PlanResult, PlanState if TYPE_CHECKING: from typing import Any @@ -303,9 +303,9 @@ class CuroboAutoGenCfg: ``None`` (default) uses ``$XDG_CACHE_HOME/embodichain_curobo`` or ``~/.cache/embodichain_curobo``. The cache key hashes the generator version, - URDF path, URDF content, control part, tool frame, and fit parameters, so - editing the URDF, changing the fit settings, or a generator update - regenerates automatically. + URDF path, URDF content, control part, tool frame, fit parameters, and initial + values of locked joints. Changing their static collision geometry or a + generator update regenerates automatically. """ fit_type: str = "voxel" @@ -444,7 +444,7 @@ class CuroboPlannerCfg(BasePlannerCfg): """Whether callers must retain cuRobo's raw collision-checked samples exactly. When ``False`` (default), - :class:`~embodichain.lab.sim.planners.motion_generator.MotionGenerator` + :class:`~embodichain.lab.sim.motion.planners.motion_generator.MotionGenerator` resamples the returned trajectory to ``MotionGenOptions.sample_count`` - matching the documented contract of :attr:`~embodichain.lab.sim.atomic_actions.MotionPolicy.sample_count` @@ -473,7 +473,7 @@ class CuroboPlanOptions(PlanOptions): """Per-plan options for :class:`CuroboPlanner`. ``start_qpos`` and ``control_part`` are populated from the - :class:`~embodichain.lab.sim.planners.motion_generator.MotionGenOptions` + :class:`~embodichain.lab.sim.motion.planners.motion_generator.MotionGenOptions` runtime context via :meth:`CuroboPlanner.with_motion_context`. """ @@ -856,9 +856,10 @@ def validate_joint_trajectory( The samples are not replanned or replaced. They are mapped from the simulator's control-part order into the exact cuRobo model, then checked against joint bounds, self-collision, and the live world collision - checker. Calling the configuration validator once per horizon sample - works around cuRobo 0.8's configuration-only ``validate`` contract while - retaining batched environments. + checker. The cuRobo 0.8 convenience validators pass obsolete tensor + arguments to the scene cost and flatten trajectory dimensions. Calling + its bound, self, and scene costs with the actual kinematics state avoids + both issues while retaining batched environments. """ if ( not isinstance(trajectory, torch.Tensor) @@ -934,13 +935,42 @@ def validate_joint_trajectory( else nullcontext() ) with device_context: + checker = backend.collision_checker for sample_index in range(horizon): sample = curobo_trajectory[:, sample_index : sample_index + 1] - valid = backend.collision_checker.validate( - sample, - env_query_idx=env_query_idx, + state = checker.get_kinematics(sample) + if sample_index == 0: + checker.setup_batch_tensors(batch_size, 1) + if checker.self_collision_cost is None: + raise RuntimeError("cuRobo self-collision cost is unavailable.") + if checker.collision_constraint is not None: + checker.collision_constraint.update_num_spheres( + state.robot_spheres.shape[-2], batch_size, 1 + ) + elif self.cfg.world.rigid_objects: + raise RuntimeError( + "cuRobo scene-collision cost is unavailable." + ) + costs = [ + checker.get_bound(sample), + checker.get_self_collision(state.robot_spheres), + ] + if checker.collision_constraint is not None: + costs.append( + checker.collision_constraint.forward( + state, idxs_env_query=env_query_idx + ) + ) + valid = torch.ones( + batch_size, device=self._curobo_device, dtype=torch.bool ) - samples.append(valid[:, 0].to(torch.bool)) + for cost in costs: + valid &= ( + (torch.isfinite(cost) & (cost == 0)) + .reshape(batch_size, -1) + .all(dim=1) + ) + samples.append(valid) return torch.stack(samples, dim=1).to(trajectory.device) def __init__(self, cfg: CuroboPlannerCfg) -> None: @@ -1268,11 +1298,18 @@ def _get_backend( batch_size: int, planning_mode: MoveType = MoveType.EEF_MOVE, ) -> "_CuroboBackend": - """Return a cached in-process backend for one goal-buffer shape.""" + """Reuse a backend only while its locked-joint model remains unchanged.""" multi_env = bool(self.cfg.world.multi_env) key = (control_part, int(batch_size), multi_env, planning_mode) + lock_signature = self._locked_joint_signature(control_part) if key in self._backend_cache: - return self._backend_cache[key] + cached = self._backend_cache[key] + if cached.robot_lock_signature != lock_signature: + raise RuntimeError( + "cuRobo locked-joint configuration changed after backend " + "creation. Call planner.close() before rebuilding its model." + ) + return cached profile = self._materialize_profile(control_part) sim_joint_names = self._resolve_sim_joint_names(control_part) @@ -1366,6 +1403,7 @@ def _get_backend( f"cuRobo capture coordinator release failed: {exc}" ) + backend.robot_lock_signature = lock_signature self._backend_cache[key] = backend logger.log_info( f"cuRobo in-process backend ready for '{control_part}' " @@ -1741,7 +1779,7 @@ def _robot_yaml_cache_key( tool_frame: str | None, auto: CuroboAutoGenCfg, ) -> str: - """Hash the URDF path/content and fit parameters into a stable cache key.""" + """Hash geometry, fit settings, and locked joint values into a cache key.""" hasher = hashlib.md5() hasher.update(_CUROBO_ROBOT_YAML_GENERATOR_VERSION.encode("utf-8")) hasher.update(urdf_path.encode("utf-8")) @@ -1758,8 +1796,30 @@ def _robot_yaml_cache_key( hasher.update(str(auto.surface_radius).encode("utf-8")) hasher.update(str(auto.iterations).encode("utf-8")) hasher.update(str(auto.collision_sphere_buffer).encode("utf-8")) + hasher.update(self._locked_joint_signature(control_part).encode("utf-8")) return hasher.hexdigest() + def _locked_joint_signature(self, control_part: str) -> str: + """Identify the static joint values shared by disk and runtime models.""" + # Match generate_curobo_robot_yaml's initial-state resolution. Joints + # outside the control part become static collision geometry, so reusing + # a YAML with different gripper/other-arm values is unsafe. + joint_names = tuple(self.robot.joint_names) + initial = self.robot.cfg.init_qpos + initial = [] if initial is None else list(initial) + if len(initial) != len(joint_names): + try: + initial = self.robot.get_qpos()[0].detach().cpu().tolist() + except Exception: + initial = [0.0] * len(joint_names) + controlled = tuple((self.robot.control_parts or {}).get(control_part, ())) + locked = tuple( + (name, float(value)) + for name, value in zip(joint_names, initial) + if name not in controlled + ) + return repr((controlled, locked)) + def _auto_generate_world_yaml(self, world_cfg: CuroboWorldCfg) -> str: """Return a cached cuRobo world YAML path generated from ``rigid_objects``. @@ -2528,6 +2588,7 @@ class _CuroboBackend: use_cuda_graph: bool planning_mode: MoveType collision_checker: "Any | None" = None + robot_lock_signature: str | None = None # Lazily-built device-tensor caches for the shared post-processing. The # cuRobo joint order and the profile's fixed transforms are stable for a # planner's life, so these are built once on first use and reused across diff --git a/embodichain/lab/sim/planners/curobo/curobo_yaml.py b/embodichain/lab/sim/motion/planners/curobo/curobo_yaml.py similarity index 99% rename from embodichain/lab/sim/planners/curobo/curobo_yaml.py rename to embodichain/lab/sim/motion/planners/curobo/curobo_yaml.py index 65015695e..55029e6f8 100644 --- a/embodichain/lab/sim/planners/curobo/curobo_yaml.py +++ b/embodichain/lab/sim/motion/planners/curobo/curobo_yaml.py @@ -20,7 +20,7 @@ spheres to every link mesh with cuRobo's sphere-fitting library, and writes a complete cuRobo V2 robot configuration YAML. The cuRobo planner adapter calls this automatically (with on-disk caching) on the first plan; see -:class:`~embodichain.lab.sim.planners.curobo.curobo_planner.CuroboAutoGenCfg`. +:class:`~embodichain.lab.sim.motion.planners.curobo.curobo_planner.CuroboAutoGenCfg`. :func:`generate_curobo_world_yaml` builds the cuRobo collision-world YAML from live :class:`~embodichain.lab.sim.objects.RigidObject` meshes. @@ -578,7 +578,7 @@ def generate_curobo_world_yaml( the simulator world origin, rebase the object poses first, or register the obstacle name in ``CuroboWorldCfg.dynamic_obstacle_names`` and update its pose at plan time via - :meth:`~embodichain.lab.sim.planners.curobo.curobo_planner.CuroboPlanner.update_dynamic_obstacles`. + :meth:`~embodichain.lab.sim.motion.planners.curobo.curobo_planner.CuroboPlanner.update_dynamic_obstacles`. Args: rigid_objects: Objects to bake into the collision world. Mapping keys are diff --git a/embodichain/lab/sim/planners/motion_generator.py b/embodichain/lab/sim/motion/planners/motion_generator.py similarity index 99% rename from embodichain/lab/sim/planners/motion_generator.py rename to embodichain/lab/sim/motion/planners/motion_generator.py index a15c224b8..c73ffa93b 100644 --- a/embodichain/lab/sim/planners/motion_generator.py +++ b/embodichain/lab/sim/motion/planners/motion_generator.py @@ -26,7 +26,7 @@ import numpy as np import torch -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( BasePlannerCfg, CollisionWorldInfo, PlanOptions, diff --git a/embodichain/lab/sim/planners/neural_planner.py b/embodichain/lab/sim/motion/planners/neural_planner.py similarity index 99% rename from embodichain/lab/sim/planners/neural_planner.py rename to embodichain/lab/sim/motion/planners/neural_planner.py index f5c3d8c90..8de69684e 100644 --- a/embodichain/lab/sim/planners/neural_planner.py +++ b/embodichain/lab/sim/motion/planners/neural_planner.py @@ -22,14 +22,14 @@ import torch import torch.nn as nn -from embodichain.lab.sim.planners.base_planner import ( +from embodichain.lab.sim.motion.planners.base_planner import ( BasePlanner, BasePlannerCfg, PlanOptions, _infer_batch_size, validate_plan_options, ) -from embodichain.lab.sim.planners.utils import MoveType, PlanResult, PlanState +from embodichain.lab.sim.motion.planners.utils import MoveType, PlanResult, PlanState from embodichain.utils import configclass, logger from embodichain.utils.math import convert_quat, quat_error_magnitude, quat_from_matrix diff --git a/embodichain/lab/sim/planners/toppra_planner.py b/embodichain/lab/sim/motion/planners/toppra_planner.py similarity index 99% rename from embodichain/lab/sim/planners/toppra_planner.py rename to embodichain/lab/sim/motion/planners/toppra_planner.py index 5d612348c..1b7783fc1 100644 --- a/embodichain/lab/sim/planners/toppra_planner.py +++ b/embodichain/lab/sim/motion/planners/toppra_planner.py @@ -23,8 +23,8 @@ from concurrent.futures.process import BrokenProcessPool from embodichain.utils import logger, configclass -from embodichain.lab.sim.planners.utils import MoveType, TrajectorySampleMethod -from embodichain.lab.sim.planners.base_planner import ( +from embodichain.lab.sim.motion.planners.utils import MoveType, TrajectorySampleMethod +from embodichain.lab.sim.motion.planners.base_planner import ( validate_plan_options, BasePlanner, BasePlannerCfg, diff --git a/embodichain/lab/sim/planners/utils.py b/embodichain/lab/sim/motion/planners/utils.py similarity index 100% rename from embodichain/lab/sim/planners/utils.py rename to embodichain/lab/sim/motion/planners/utils.py diff --git a/embodichain/lab/sim/solvers/__init__.py b/embodichain/lab/sim/motion/solvers/__init__.py similarity index 77% rename from embodichain/lab/sim/solvers/__init__.py rename to embodichain/lab/sim/motion/solvers/__init__.py index 70a7c386f..423a26a20 100644 --- a/embodichain/lab/sim/solvers/__init__.py +++ b/embodichain/lab/sim/motion/solvers/__init__.py @@ -19,6 +19,8 @@ Provides analytic (SRS, OPW, UR), numerical (Pinocchio, Pink, Differential), and learning-based (PyTorch, NeuralIK) solvers, one per robot control part. """ +from __future__ import annotations + from .base_solver import SolverCfg, BaseSolver from .pytorch_solver import PytorchSolverCfg, PytorchSolver from .pinocchio_solver import PinocchioSolverCfg, PinocchioSolver @@ -28,3 +30,24 @@ from .srs_solver import SRSSolverCfg, SRSSolver from .neural_ik_solver import NeuralIKSolverCfg, NeuralIKSolver from .ur_solver import URSolverCfg, URSolver + +__all__ = [ + "SolverCfg", + "BaseSolver", + "PytorchSolverCfg", + "PytorchSolver", + "PinocchioSolverCfg", + "PinocchioSolver", + "DifferentialSolverCfg", + "DifferentialSolver", + "PinkSolverCfg", + "PinkSolver", + "OPWSolverCfg", + "OPWSolver", + "SRSSolverCfg", + "SRSSolver", + "NeuralIKSolverCfg", + "NeuralIKSolver", + "URSolverCfg", + "URSolver", +] diff --git a/embodichain/lab/sim/solvers/base_solver.py b/embodichain/lab/sim/motion/solvers/base_solver.py similarity index 99% rename from embodichain/lab/sim/solvers/base_solver.py rename to embodichain/lab/sim/motion/solvers/base_solver.py index a9bb9e1a5..e3844aa6e 100644 --- a/embodichain/lab/sim/solvers/base_solver.py +++ b/embodichain/lab/sim/motion/solvers/base_solver.py @@ -115,7 +115,7 @@ def from_dict(cls, init_dict: Dict[str, Any]) -> "SolverCfg": logger.log_error("class type must be specified in the configuration.") cfg_type = get_class_instance( - "embodichain.lab.sim.solvers", init_dict["class_type"] + "Cfg" + "embodichain.lab.sim.motion.solvers", init_dict["class_type"] + "Cfg" ) concrete_fields = {field.name: field for field in fields(cfg_type)} kwargs: Dict[str, Any] = {} diff --git a/embodichain/lab/sim/solvers/differential_solver.py b/embodichain/lab/sim/motion/solvers/differential_solver.py similarity index 99% rename from embodichain/lab/sim/solvers/differential_solver.py rename to embodichain/lab/sim/motion/solvers/differential_solver.py index 12e51bcbd..f3925be4b 100644 --- a/embodichain/lab/sim/solvers/differential_solver.py +++ b/embodichain/lab/sim/motion/solvers/differential_solver.py @@ -19,7 +19,7 @@ from scipy.spatial.transform import Rotation from embodichain.utils import configclass, logger -from embodichain.lab.sim.solvers import SolverCfg, BaseSolver +from embodichain.lab.sim.motion.solvers import SolverCfg, BaseSolver from embodichain.utils.math import ( apply_delta_pose, compute_pose_error, diff --git a/embodichain/lab/sim/solvers/neural_ik_solver.py b/embodichain/lab/sim/motion/solvers/neural_ik_solver.py similarity index 98% rename from embodichain/lab/sim/solvers/neural_ik_solver.py rename to embodichain/lab/sim/motion/solvers/neural_ik_solver.py index 7f1cb1d12..56159df1a 100644 --- a/embodichain/lab/sim/solvers/neural_ik_solver.py +++ b/embodichain/lab/sim/motion/solvers/neural_ik_solver.py @@ -24,8 +24,8 @@ quat_error_magnitude, quat_from_matrix, ) -from embodichain.lab.sim.solvers import SolverCfg, BaseSolver -from embodichain.lab.sim.solvers.qpos_seed_sampler import QposSeedSampler +from embodichain.lab.sim.motion.solvers import SolverCfg, BaseSolver +from embodichain.lab.sim.motion.solvers.qpos_seed_sampler import QposSeedSampler __all__ = ["NeuralIKSolverCfg", "NeuralIKSolver"] diff --git a/embodichain/lab/sim/solvers/null_space_posture_task.py b/embodichain/lab/sim/motion/solvers/null_space_posture_task.py similarity index 100% rename from embodichain/lab/sim/solvers/null_space_posture_task.py rename to embodichain/lab/sim/motion/solvers/null_space_posture_task.py diff --git a/embodichain/lab/sim/solvers/opw_solver.py b/embodichain/lab/sim/motion/solvers/opw_solver.py similarity index 99% rename from embodichain/lab/sim/solvers/opw_solver.py rename to embodichain/lab/sim/motion/solvers/opw_solver.py index e3202597e..acf5bf7d8 100644 --- a/embodichain/lab/sim/solvers/opw_solver.py +++ b/embodichain/lab/sim/motion/solvers/opw_solver.py @@ -24,7 +24,7 @@ from scipy.spatial.transform import Rotation from embodichain.utils import configclass, logger -from embodichain.lab.sim.solvers import SolverCfg, BaseSolver +from embodichain.lab.sim.motion.solvers import SolverCfg, BaseSolver from embodichain.utils.warp.kinematics.opw_solver import ( OPWparam, opw_fk_kernel, diff --git a/embodichain/lab/sim/solvers/pink_solver.py b/embodichain/lab/sim/motion/solvers/pink_solver.py similarity index 98% rename from embodichain/lab/sim/solvers/pink_solver.py rename to embodichain/lab/sim/motion/solvers/pink_solver.py index d99d06f60..c342a5cb0 100644 --- a/embodichain/lab/sim/solvers/pink_solver.py +++ b/embodichain/lab/sim/motion/solvers/pink_solver.py @@ -25,7 +25,7 @@ import numpy as np import torch -from embodichain.lab.sim.solvers import BaseSolver, SolverCfg +from embodichain.lab.sim.motion.solvers import BaseSolver, SolverCfg from embodichain.lab.sim.utility.import_utils import ( lazy_import_pinocchio, lazy_import_pink, @@ -136,7 +136,7 @@ def __init__(self, cfg: PinkSolverCfg, **kwargs: Any) -> None: self.pin = lazy_import_pinocchio() self.pink = lazy_import_pink() - from embodichain.lab.sim.solvers.null_space_posture_task import ( + from embodichain.lab.sim.motion.solvers.null_space_posture_task import ( NullSpacePostureTask, ) @@ -436,7 +436,7 @@ def update_null_space_joint_targets( Args: current_qpos: Joint target in simulator ordering. """ - from embodichain.lab.sim.solvers.null_space_posture_task import ( + from embodichain.lab.sim.motion.solvers.null_space_posture_task import ( NullSpacePostureTask, ) @@ -530,7 +530,7 @@ def _set_target(self, target_xpos: np.ndarray) -> None: def _task_metrics(self) -> tuple[float, float, float, float]: """Return lexicographic task merits and frame convergence errors.""" - from embodichain.lab.sim.solvers.null_space_posture_task import ( + from embodichain.lab.sim.motion.solvers.null_space_posture_task import ( NullSpacePostureTask, ) diff --git a/embodichain/lab/sim/solvers/pinocchio_solver.py b/embodichain/lab/sim/motion/solvers/pinocchio_solver.py similarity index 99% rename from embodichain/lab/sim/solvers/pinocchio_solver.py rename to embodichain/lab/sim/motion/solvers/pinocchio_solver.py index 9ddde65b6..991ed2320 100644 --- a/embodichain/lab/sim/solvers/pinocchio_solver.py +++ b/embodichain/lab/sim/motion/solvers/pinocchio_solver.py @@ -22,7 +22,7 @@ from copy import deepcopy from embodichain.utils import configclass, logger -from embodichain.lab.sim.solvers import SolverCfg, BaseSolver +from embodichain.lab.sim.motion.solvers import SolverCfg, BaseSolver from embodichain.lab.sim.utility.import_utils import ( lazy_import_pinocchio, diff --git a/embodichain/lab/sim/solvers/pytorch_solver.py b/embodichain/lab/sim/motion/solvers/pytorch_solver.py similarity index 99% rename from embodichain/lab/sim/solvers/pytorch_solver.py rename to embodichain/lab/sim/motion/solvers/pytorch_solver.py index 277eaf877..d1d257351 100644 --- a/embodichain/lab/sim/solvers/pytorch_solver.py +++ b/embodichain/lab/sim/motion/solvers/pytorch_solver.py @@ -21,8 +21,8 @@ from copy import deepcopy from embodichain.utils import configclass, logger -from embodichain.lab.sim.solvers import SolverCfg, BaseSolver -from embodichain.lab.sim.solvers.qpos_seed_sampler import QposSeedSampler +from embodichain.lab.sim.motion.solvers import SolverCfg, BaseSolver +from embodichain.lab.sim.motion.solvers.qpos_seed_sampler import QposSeedSampler from embodichain.lab.sim.utility.solver_utils import validate_iteration_params if TYPE_CHECKING: diff --git a/embodichain/lab/sim/solvers/qpos_seed_sampler.py b/embodichain/lab/sim/motion/solvers/qpos_seed_sampler.py similarity index 100% rename from embodichain/lab/sim/solvers/qpos_seed_sampler.py rename to embodichain/lab/sim/motion/solvers/qpos_seed_sampler.py diff --git a/embodichain/lab/sim/solvers/srs_solver.py b/embodichain/lab/sim/motion/solvers/srs_solver.py similarity index 99% rename from embodichain/lab/sim/solvers/srs_solver.py rename to embodichain/lab/sim/motion/solvers/srs_solver.py index aad8cce48..bb1679d66 100644 --- a/embodichain/lab/sim/solvers/srs_solver.py +++ b/embodichain/lab/sim/motion/solvers/srs_solver.py @@ -23,7 +23,7 @@ import torch import warp as wp -from embodichain.lab.sim.solvers import BaseSolver, SolverCfg +from embodichain.lab.sim.motion.solvers import BaseSolver, SolverCfg from embodichain.utils import configclass, logger from embodichain.utils.device_utils import standardize_device_string from embodichain.utils.warp.kinematics.srs_solver import ( diff --git a/embodichain/lab/sim/solvers/ur_solver.py b/embodichain/lab/sim/motion/solvers/ur_solver.py similarity index 99% rename from embodichain/lab/sim/solvers/ur_solver.py rename to embodichain/lab/sim/motion/solvers/ur_solver.py index 692f0839e..8419f8de6 100644 --- a/embodichain/lab/sim/solvers/ur_solver.py +++ b/embodichain/lab/sim/motion/solvers/ur_solver.py @@ -18,7 +18,7 @@ import numpy as np import warp as wp from embodichain.utils import configclass -from embodichain.lab.sim.solvers import SolverCfg, BaseSolver +from embodichain.lab.sim.motion.solvers import SolverCfg, BaseSolver from embodichain.data import get_data_path from embodichain.utils.warp.kinematics.ur_solver import ( URParam, diff --git a/embodichain/lab/sim/motion/trajectory_augmentation/__init__.py b/embodichain/lab/sim/motion/trajectory_augmentation/__init__.py new file mode 100644 index 000000000..f25ca80ed --- /dev/null +++ b/embodichain/lab/sim/motion/trajectory_augmentation/__init__.py @@ -0,0 +1,70 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Host-independent grasp/qpos augmentation, coverage, and generation bookkeeping. + +Execution, initial-state restoration, physical validation, and durable storage +are supplied by host integrations. These algorithms do not call Gym or a +simulation backend, but the public import follows normal ``lab.sim`` +initialization and requires the simulation package dependencies. +""" + +from __future__ import annotations + +from .cfg import TrajectoryAugmentationCfg, TrajectoryGenerationJobCfg +from .contracts import ( + CandidateIdentity, + CandidateTrajectoryBatch, + CommitReceipt, + ExpertEpisode, + MotionSnapshot, + SceneCase, + TrajectoryPhase, + TrajectoryTemplate, + ValidationCheck, + ValidationResult, +) +from .coverage import CoverageIndex, TrajectoryDescriptor, describe_trajectory +from .operators import ( + joint_residual, + retime, + rotate_grasp_about_object_axis, + validate_motion_limits, +) +from .session import GenerationSession + +__all__ = [ + "CandidateIdentity", + "CandidateTrajectoryBatch", + "CommitReceipt", + "ExpertEpisode", + "MotionSnapshot", + "SceneCase", + "TrajectoryAugmentationCfg", + "TrajectoryGenerationJobCfg", + "TrajectoryPhase", + "TrajectoryTemplate", + "ValidationCheck", + "ValidationResult", + "CoverageIndex", + "TrajectoryDescriptor", + "describe_trajectory", + "joint_residual", + "retime", + "rotate_grasp_about_object_axis", + "validate_motion_limits", + "GenerationSession", +] diff --git a/embodichain/lab/sim/motion/trajectory_augmentation/cfg.py b/embodichain/lab/sim/motion/trajectory_augmentation/cfg.py new file mode 100644 index 000000000..27218a799 --- /dev/null +++ b/embodichain/lab/sim/motion/trajectory_augmentation/cfg.py @@ -0,0 +1,436 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Strict configuration for the first fixed-scene generation implementation. + +These are pure job contracts, not a runnable host integration. A runner must +still resolve registered sources, restoration profiles, physical validators, +and a persistence sink before execution. Capability checks use explicit trusted +registries; configuration strings are never imported or evaluated. The initial +subset rejects modes that require unimplemented scheduling or retry policies. +""" + +from __future__ import annotations + +from collections.abc import Collection, Mapping +from dataclasses import fields +import math +from typing import Any, get_args, get_origin, get_type_hints + +from embodichain.utils import configclass + +__all__ = ["TrajectoryAugmentationCfg", "TrajectoryGenerationJobCfg"] + + +def _positive(value: float, name: str, *, allow_zero: bool = False) -> None: + if ( + type(value) not in (int, float) + or not math.isfinite(value) + or (value < 0 if allow_zero else value <= 0) + ): + raise ValueError( + f"{name} must be finite and {'nonnegative' if allow_zero else 'positive'}" + ) + + +def _count(value: int, name: str, minimum: int = 1) -> None: + if type(value) is not int or value < minimum: + raise ValueError(f"{name} must be an integer >= {minimum}") + + +def _boolean(value: bool, name: str) -> None: + if type(value) is not bool: + raise ValueError(f"{name} must be a boolean") + + +def _id(value: str, name: str) -> None: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{name} must be a nonempty registered ID") + + +def _fixed(value: str, supported: str, name: str) -> None: + if value != supported: + raise ValueError(f"{name} currently supports only {supported!r}") + + +def _decode(cls: type, data: Mapping[str, Any], path: str = "") -> Any: + if not isinstance(data, Mapping): + raise ValueError(f"{path or cls.__name__} must be a mapping") + hints = get_type_hints(cls) + allowed = {field.name for field in fields(cls)} + unknown = set(data) - allowed + if unknown: + raise ValueError( + f"unknown fields in {path or cls.__name__}: {sorted(unknown, key=str)}" + ) + decoded = {} + for key, value in data.items(): + expected = hints[key] + name = f"{path}.{key}" if path else key + if hasattr(expected, "__dataclass_fields__"): + decoded[key] = _decode(expected, value, name) + elif get_origin(expected) is tuple: + if not isinstance(value, (list, tuple)): + raise ValueError(f"{name} must be a sequence") + member = get_args(expected)[0] + if any( + type(item) not in ((float, int) if member is float else (member,)) + for item in value + ): + raise ValueError(f"{name} contains an invalid value type") + decoded[key] = tuple(value) + elif expected is float and type(value) in (float, int): + decoded[key] = float(value) + elif type(value) is expected: + decoded[key] = value + else: + raise ValueError(f"{name} has an invalid type") + return cls(**decoded) + + +@configclass +class _ProvidedStartCfg: + mode: str = "provided" + + def __post_init__(self) -> None: + _fixed(self.mode, "provided", "start_state.mode") + + +@configclass +class _DisabledFactorCfg: + enabled: bool = False + + def __post_init__(self) -> None: + _boolean(self.enabled, "factor.enabled") + if self.enabled: + raise ValueError("this augmentation factor is not implemented") + + +@configclass +class _SpatialCfg: + enabled: bool = False + method: str = "joint_residual" + joint_offset_scale: float = 0.05 + + def __post_init__(self) -> None: + _boolean(self.enabled, "spatial.enabled") + if self.method not in ("joint_residual", "via_points"): + raise ValueError("spatial.method must be joint_residual or via_points") + _positive( + self.joint_offset_scale, "spatial.joint_offset_scale", allow_zero=True + ) + if self.joint_offset_scale > 1: + raise ValueError( + "spatial.joint_offset_scale is normalized and must be <= 1" + ) + + +@configclass +class _TimingCfg: + enabled: bool = False + duration_scales: tuple[float, ...] = (1.0,) + + def __post_init__(self) -> None: + _boolean(self.enabled, "timing.enabled") + if ( + not isinstance(self.duration_scales, (list, tuple)) + or not self.duration_scales + ): + raise ValueError("timing.duration_scales must be a nonempty sequence") + for value in self.duration_scales: + _positive(value, "timing.duration_scales") + if len(set(self.duration_scales)) != len(self.duration_scales): + raise ValueError("timing.duration_scales must be unique") + self.duration_scales = tuple(self.duration_scales) + if not self.enabled and self.duration_scales != (1.0,): + raise ValueError("disabled timing must retain the reference duration scale") + + +@configclass +class _FactorsCfg: + contact: _DisabledFactorCfg = _DisabledFactorCfg() + ik: _DisabledFactorCfg = _DisabledFactorCfg() + approach: _DisabledFactorCfg = _DisabledFactorCfg() + spatial: _SpatialCfg = _SpatialCfg() + timing: _TimingCfg = _TimingCfg() + contact_timing: _DisabledFactorCfg = _DisabledFactorCfg() + recovery: _DisabledFactorCfg = _DisabledFactorCfg() + + +@configclass +class _CoverageCfg: + geometry_samples_per_phase: int = 32 + joint_dedup_normalized_tol: float = 0.01 + target_per_cell: int = 1 + + def __post_init__(self) -> None: + _count(self.geometry_samples_per_phase, "geometry_samples_per_phase", 2) + _positive(self.joint_dedup_normalized_tol, "joint_dedup_normalized_tol") + if self.joint_dedup_normalized_tol > 1: + raise ValueError( + "joint_dedup_normalized_tol is normalized and must be <= 1" + ) + _count(self.target_per_cell, "target_per_cell") + + +@configclass +class TrajectoryAugmentationCfg: + """Local random seed, explicitly enabled factors, and geometry coverage limits.""" + + seed: int = 0 + start_state: _ProvidedStartCfg = _ProvidedStartCfg() + factors: _FactorsCfg = _FactorsCfg() + coverage: _CoverageCfg = _CoverageCfg() + + def __post_init__(self) -> None: + _count(self.seed, "seed", 0) + if self.seed >= 2**63: + raise ValueError("seed must be less than 2**63") + self.validate_semantics() + + @classmethod + def from_mapping(cls, data: Mapping[str, Any]) -> TrajectoryAugmentationCfg: + """Decode closed nested fields without executing configuration values.""" + return _decode(cls, data) + + def validate_semantics(self) -> None: + """Recheck typed nested values, including mutations made after decoding.""" + for name, expected in ( + ("start_state", _ProvidedStartCfg), + ("factors", _FactorsCfg), + ("coverage", _CoverageCfg), + ): + _validate_nested(getattr(self, name), expected, name) + _count(self.seed, "seed", 0) + if self.seed >= 2**63: + raise ValueError("seed must be less than 2**63") + + +@configclass +class _SourceCfg: + kind: str = "handwritten" + source_id: str = "handwritten_qpos" + template_id: str = "reference_0" + + def __post_init__(self) -> None: + if self.kind not in ("handwritten", "atomic"): + raise ValueError("source.kind must be handwritten or atomic") + _id(self.source_id, "source_id") + _id(self.template_id, "template_id") + + +@configclass +class _PlanningCfg: + batch_mode: str = "env_rows" + + def __post_init__(self) -> None: + _fixed(self.batch_mode, "env_rows", "planning.batch_mode") + + +@configclass +class _ExecutionCfg: + pool_mode: str = "per_env_case" + scheduler: str = "full_batch" + ready_low_watermark: int = 1 + ready_high_watermark: int = 16 + ready_max_bytes: int = 268435456 + overlap_planning_and_physics: bool = False + + def __post_init__(self) -> None: + _fixed(self.pool_mode, "per_env_case", "execution.pool_mode") + _fixed(self.scheduler, "full_batch", "execution.scheduler") + _count(self.ready_low_watermark, "ready_low_watermark", 0) + _count(self.ready_high_watermark, "ready_high_watermark") + _count(self.ready_max_bytes, "ready_max_bytes") + if self.ready_low_watermark >= self.ready_high_watermark: + raise ValueError( + "ready_low_watermark must be less than ready_high_watermark" + ) + _boolean(self.overlap_planning_and_physics, "overlap_planning_and_physics") + if self.overlap_planning_and_physics: + raise ValueError("overlapping planning and physics is not implemented") + + +@configclass +class _ResetCfg: + outer_mode: str = "provided" + inner_mode: str = "restore_initial" + prepare_profile_id: str = "fixed_scene_initial_state" + initial_state_tolerances_profile_id: str = "fixed_scene_tolerances" + on_initial_state_mismatch: str = "error" + + def __post_init__(self) -> None: + _fixed(self.outer_mode, "provided", "reset.outer_mode") + _fixed(self.inner_mode, "restore_initial", "reset.inner_mode") + _fixed(self.on_initial_state_mismatch, "error", "on_initial_state_mismatch") + _id(self.prepare_profile_id, "prepare_profile_id") + _id( + self.initial_state_tolerances_profile_id, + "initial_state_tolerances_profile_id", + ) + + +@configclass +class _ValidationCfg: + validator_id: str = "task_success" + profile_id: str = "verified_motion" + motion_limits_profile_id: str = "robot_execution_limits" + require_path_collision: bool = True + require_task_success: bool = True + path_length_ratio_max: float = 1.5 + duration_ratio_max: float = 1.5 + + def __post_init__(self) -> None: + for name in ("validator_id", "profile_id", "motion_limits_profile_id"): + _id(getattr(self, name), name) + for name in ("require_path_collision", "require_task_success"): + _boolean(getattr(self, name), name) + if not getattr(self, name): + raise ValueError(f"expert collection requires {name}") + _positive(self.path_length_ratio_max, "path_length_ratio_max") + _positive(self.duration_ratio_max, "duration_ratio_max") + + +@configclass +class _CollectionCfg: + target_committed_episodes: int = 1 + max_proposals: int = 100 + max_rollout_attempts: int = 100 + max_attempts_per_candidate: int = 1 + max_wall_time_s: float = 60.0 + + def __post_init__(self) -> None: + for name in ( + "target_committed_episodes", + "max_proposals", + "max_rollout_attempts", + "max_attempts_per_candidate", + ): + _count(getattr(self, name), name) + _positive(self.max_wall_time_s, "max_wall_time_s") + if self.max_attempts_per_candidate != 1: + raise ValueError( + "rollout retries are not implemented; max_attempts_per_candidate must be 1" + ) + if self.target_committed_episodes > min( + self.max_proposals, self.max_rollout_attempts + ): + raise ValueError( + "collection budgets cannot be smaller than the committed target" + ) + + +@configclass +class _PersistenceCfg: + sink: str = "lerobot" + accepted_only: bool = True + async_write: bool = False + pending_episode_limit: int = 8 + pending_max_bytes: int = 536870912 + save_audit: bool = True + split_unit: str = "trajectory_family" + + def __post_init__(self) -> None: + _id(self.sink, "sink") + for name in ("accepted_only", "async_write", "save_audit"): + _boolean(getattr(self, name), name) + if not self.accepted_only or not self.save_audit: + raise ValueError("expert collection requires accepted_only and save_audit") + if self.async_write: + raise ValueError("asynchronous persistence is not implemented") + _count(self.pending_episode_limit, "pending_episode_limit") + _count(self.pending_max_bytes, "pending_max_bytes") + _fixed(self.split_unit, "trajectory_family", "split_unit") + + +def _validate_nested(value: Any, expected: type, name: str) -> None: + if not isinstance(value, expected): + raise ValueError( + f"{name} must use its typed configuration; use from_mapping for dictionaries" + ) + # Re-decode a plain mapping to reject mutated nested types and unsupported values. + _decode(expected, value.to_dict(), name) + + +@configclass +class TrajectoryGenerationJobCfg: + """Standalone generation job configuration; control periods remain host-owned.""" + + source: _SourceCfg = _SourceCfg() + augmentation: TrajectoryAugmentationCfg = TrajectoryAugmentationCfg() + planning: _PlanningCfg = _PlanningCfg() + execution: _ExecutionCfg = _ExecutionCfg() + reset: _ResetCfg = _ResetCfg() + validation: _ValidationCfg = _ValidationCfg() + collection: _CollectionCfg = _CollectionCfg() + persistence: _PersistenceCfg = _PersistenceCfg() + + def __post_init__(self) -> None: + self.validate_semantics() + + @classmethod + def from_mapping(cls, data: Mapping[str, Any]) -> TrajectoryGenerationJobCfg: + """Decode the supported nested YAML schema and reject unknown fields.""" + return _decode(cls, data) + + def validate_semantics(self) -> None: + """Validate configuration without loading hosts, profiles, or sources.""" + for name, expected in get_type_hints(type(self)).items(): + _validate_nested(getattr(self, name), expected, name) + + def validate_capabilities( + self, + *, + source_ids: Collection[str], + validator_ids: Collection[str], + profile_ids: Collection[str], + sink_ids: Collection[str], + operators: Collection[str] = (), + ) -> None: + """Require IDs and enabled operators in explicitly supplied registries. + + This is a preflight step in addition to semantic decoding. Registry + membership asserts that the caller has already resolved trusted + implementations; it does not itself certify physical validation. + """ + self.validate_semantics() + references = ( + (self.source.source_id, source_ids, "source"), + (self.validation.validator_id, validator_ids, "validator"), + (self.persistence.sink, sink_ids, "sink"), + (self.reset.prepare_profile_id, profile_ids, "prepare profile"), + ( + self.reset.initial_state_tolerances_profile_id, + profile_ids, + "initial-state profile", + ), + (self.validation.profile_id, profile_ids, "validation profile"), + ( + self.validation.motion_limits_profile_id, + profile_ids, + "motion-limits profile", + ), + ) + for value, registry, name in references: + if value not in registry: + raise ValueError(f"unregistered {name}: {value!r}") + spatial = self.augmentation.factors.spatial + if spatial.enabled and spatial.method not in operators: + raise ValueError( + f"spatial operator capability unavailable: {spatial.method!r}" + ) + if self.augmentation.factors.timing.enabled and "retime" not in operators: + raise ValueError("retime operator capability unavailable") diff --git a/embodichain/lab/sim/motion/trajectory_augmentation/contracts.py b/embodichain/lab/sim/motion/trajectory_augmentation/contracts.py new file mode 100644 index 000000000..649ebdd80 --- /dev/null +++ b/embodichain/lab/sim/motion/trajectory_augmentation/contracts.py @@ -0,0 +1,582 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Owned values for fixed-scene trajectory generation, independent of hosts. + +Tensor inputs are detached and copied at construction. Frozen dataclasses prevent +field replacement; consumers must still treat their owned tensors as read-only. +Positions use the complete declared joint order. All times are in seconds. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +import math +from types import MappingProxyType +from typing import Literal + +import torch + +__all__ = [ + "SceneCase", + "MotionSnapshot", + "TrajectoryPhase", + "TrajectoryTemplate", + "CandidateIdentity", + "CandidateTrajectoryBatch", + "ValidationCheck", + "ValidationResult", + "ExpertEpisode", + "CommitReceipt", +] + + +def _text(value: str, name: str) -> None: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{name} must be a nonempty string") + + +def _integer(value: int, name: str, minimum: int = 0) -> None: + if type(value) is not int or value < minimum: + raise ValueError(f"{name} must be an integer >= {minimum}") + + +def _names(values: tuple[str, ...], name: str) -> tuple[str, ...]: + if isinstance(values, str): + raise ValueError(f"{name} must be a sequence of names") + result = tuple(values) + for value in result: + _text(value, name) + if len(set(result)) != len(result): + raise ValueError(f"{name} must be unique") + return result + + +def _tensor(value: torch.Tensor, name: str, ndim: int) -> torch.Tensor: + if not isinstance(value, torch.Tensor) or value.ndim != ndim: + raise ValueError(f"{name} must be a rank-{ndim} tensor") + if not value.is_floating_point() or not bool(torch.isfinite(value).all()): + raise ValueError(f"{name} must contain finite floating point values") + return value.detach().clone() + + +def _pose(value: torch.Tensor, name: str) -> torch.Tensor: + result = _tensor(value, name, 2) + if result.shape != (4, 4): + raise ValueError(f"{name} must have shape (4, 4)") + # Float32 transforms accumulated by a host can carry small roundoff errors. + check = result.to(dtype=torch.float64) + rotation = check[:3, :3] + identity = torch.eye(3, dtype=check.dtype, device=check.device) + bottom = check.new_tensor([0, 0, 0, 1]) + if ( + not torch.allclose(check[3], bottom, atol=1e-5, rtol=0) + or not torch.allclose(rotation.T @ rotation, identity, atol=1e-5, rtol=0) + or not torch.allclose( + torch.linalg.det(rotation), check.new_tensor(1), atol=1e-5, rtol=0 + ) + ): + raise ValueError(f"{name} must be a proper SE(3) transform") + return result + + +def _poses(values: Mapping[str, torch.Tensor], name: str) -> Mapping[str, torch.Tensor]: + result = {} + for key, value in values.items(): + _text(key, name) + result[key] = _pose(value, f"{name}[{key}]") + return MappingProxyType(result) + + +def _json(value: object) -> object: + if value is None or type(value) in (str, bool, int): + return value + if type(value) is float and math.isfinite(value): + return value + if isinstance(value, (list, tuple)): + return tuple(_json(item) for item in value) + if isinstance(value, Mapping): + result = {} + for key, item in value.items(): + _text(key, "metadata key") + result[key] = _json(item) + return MappingProxyType(result) + raise ValueError("metadata must contain only finite JSON values") + + +def _phases( + values: tuple[TrajectoryPhase, ...], length: int +) -> tuple[TrajectoryPhase, ...]: + result = tuple(values) + previous_stop = 0 + seen = set() + for phase in result: + if not isinstance(phase, TrajectoryPhase): + raise ValueError("phases must contain TrajectoryPhase values") + if ( + phase.phase_id in seen + or phase.start_index < previous_stop + or phase.stop_index > length + ): + raise ValueError( + "phases must be unique, ordered, nonoverlapping, and within valid samples" + ) + seen.add(phase.phase_id) + previous_stop = phase.stop_index + return result + + +@dataclass(frozen=True) +class SceneCase: + """A fixed scene and allowed initial state, with no live environment references.""" + + scene_case_id: str + initial_state_id: str + scene_signature: str + task_id: str + robot_profile_id: str + calibration_id: str = "default" + + def __post_init__(self) -> None: + for name in self.__dataclass_fields__: + _text(getattr(self, name), name) + + +@dataclass(frozen=True) +class MotionSnapshot: + """One copied complete robot state; poses map local frames into the world.""" + + scene_case: SceneCase + joint_names: tuple[str, ...] + joint_positions: torch.Tensor + joint_velocities: torch.Tensor + root_pose: torch.Tensor + entity_poses: Mapping[str, torch.Tensor] = field(default_factory=dict) + dependency_revisions: Mapping[str, int] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.scene_case, SceneCase): + raise ValueError("scene_case must be a SceneCase") + names = _names(self.joint_names, "joint_names") + positions = _tensor(self.joint_positions, "joint_positions", 1) + velocities = _tensor(self.joint_velocities, "joint_velocities", 1) + if ( + not names + or positions.shape != (len(names),) + or velocities.shape != positions.shape + ): + raise ValueError( + "robot positions and velocities must match the complete joint_names" + ) + if velocities.device != positions.device: + raise ValueError("robot positions and velocities must share a device") + revisions = dict(self.dependency_revisions) + for key, revision in revisions.items(): + _text(key, "dependency revision key") + _integer(revision, "dependency revision") + object.__setattr__(self, "joint_names", names) + object.__setattr__(self, "joint_positions", positions) + object.__setattr__(self, "joint_velocities", velocities) + object.__setattr__(self, "root_pose", _pose(self.root_pose, "root_pose")) + object.__setattr__( + self, "entity_poses", _poses(self.entity_poses, "entity_poses") + ) + object.__setattr__(self, "dependency_revisions", MappingProxyType(revisions)) + + +@dataclass(frozen=True) +class TrajectoryPhase: + """An explicitly annotated half-open sample interval ``[start, stop)``. + + Phase endpoints are samples ``start_index`` and ``stop_index - 1``. + Unannotated samples remain unchanged by augmentation operators. + """ + + phase_id: str + start_index: int + stop_index: int + kind: Literal["free", "contact", "hold"] = "free" + allowed_operators: tuple[str, ...] = () + + def __post_init__(self) -> None: + _text(self.phase_id, "phase_id") + _integer(self.start_index, "start_index") + _integer(self.stop_index, "stop_index", self.start_index + 1) + if self.kind not in ("free", "contact", "hold"): + raise ValueError("phase kind must be free, contact, or hold") + object.__setattr__( + self, + "allowed_operators", + _names(self.allowed_operators, "allowed_operators"), + ) + + +@dataclass(frozen=True) +class TrajectoryTemplate: + """A full-joint qpos reference with arrival intervals and explicit permissions. + + ``dt[0]`` is zero; every later interval is strictly positive. Empty phases + and operator permissions permit replay only. EEF conversion is owned by + an injected planning adapter before constructing this qpos contract. + """ + + source_id: str + source_revision: str + template_id: str + joint_names: tuple[str, ...] + positions: torch.Tensor + dt: torch.Tensor + phases: tuple[TrajectoryPhase, ...] = () + allowed_operators: tuple[str, ...] = () + validator_id: str = "default" + controlled_joint_indices: tuple[int, ...] = () + representation: str = "qpos" + + def __post_init__(self) -> None: + for name in ("source_id", "source_revision", "template_id", "validator_id"): + _text(getattr(self, name), name) + if self.representation != "qpos": + raise ValueError("only explicit qpos templates are supported") + names = _names(self.joint_names, "joint_names") + positions = _tensor(self.positions, "positions", 2) + dt = _tensor(self.dt, "dt", 1) + if not names or positions.shape[0] < 1 or positions.shape[1] != len(names): + raise ValueError("positions must have shape (N >= 1, len(joint_names))") + if dt.shape != positions.shape[:1] or dt.device != positions.device: + raise ValueError("dt must match the positions length and device") + if dt[0] != 0 or bool((dt[1:] <= 0).any()): + raise ValueError("dt[0] must be zero and later arrival intervals positive") + controlled = tuple(self.controlled_joint_indices) + for index in controlled: + _integer(index, "controlled joint index") + if index >= len(names): + raise ValueError("controlled joint index is outside joint_names") + if len(set(controlled)) != len(controlled): + raise ValueError("controlled joint indices must be unique") + object.__setattr__(self, "joint_names", names) + object.__setattr__(self, "positions", positions) + object.__setattr__(self, "dt", dt) + object.__setattr__(self, "phases", _phases(self.phases, positions.shape[0])) + object.__setattr__( + self, + "allowed_operators", + _names(self.allowed_operators, "allowed_operators"), + ) + object.__setattr__(self, "controlled_joint_indices", controlled) + + +@dataclass(frozen=True) +class CandidateIdentity: + """Candidate lineage independent of execution slots and their epochs.""" + + scene_case_id: str + initial_state_id: str + candidate_id: str + geometry_family_id: str + source_id: str + source_revision: str + template_id: str + attempt_id: int = 0 + parent_id: str | None = None + + def __post_init__(self) -> None: + for name in ( + "scene_case_id", + "initial_state_id", + "candidate_id", + "geometry_family_id", + "source_id", + "source_revision", + "template_id", + ): + _text(getattr(self, name), name) + _integer(self.attempt_id, "attempt_id") + if self.parent_id is not None: + _text(self.parent_id, "parent_id") + if self.parent_id == self.candidate_id: + raise ValueError("a candidate cannot be its own parent") + + +@dataclass(frozen=True) +class CandidateTrajectoryBatch: + """Logical candidate rows with safe padding and no physical batch assumption. + + ``positions`` is ``(C,N,D_full)``, ``dt`` is ``(C,N)``, and each + ``valid_length`` is in ``[1,N]``. An empty candidate set uses ``C=0``. + Padded positions hold the final valid sample and padded intervals are zero. + They must never be executed or recorded; use :attr:`valid_mask`. + """ + + positions: torch.Tensor + dt: torch.Tensor + valid_length: torch.Tensor + identities: tuple[CandidateIdentity, ...] + joint_names: tuple[str, ...] + phases: tuple[tuple[TrajectoryPhase, ...], ...] = () + factors: tuple[Mapping[str, float | int | str], ...] = () + source_row_indices: torch.Tensor | None = None + + def __post_init__(self) -> None: + positions = _tensor(self.positions, "positions", 3) + dt = _tensor(self.dt, "dt", 2) + count, horizon, joints = positions.shape + names = _names(self.joint_names, "joint_names") + if not names or joints != len(names) or (count and horizon < 1): + raise ValueError( + "positions must use complete joint_names and a nonempty horizon" + ) + if dt.shape != (count, horizon) or dt.device != positions.device: + raise ValueError("dt must match positions shape and device") + lengths = self.valid_length + if ( + not isinstance(lengths, torch.Tensor) + or lengths.shape != (count,) + or lengths.dtype != torch.int64 + ): + raise ValueError("valid_length must be an int64 tensor of shape (C,)") + lengths = lengths.detach().clone().to(device=positions.device) + if bool(((lengths < 1) | (lengths > horizon)).any()): + raise ValueError("valid_length must be between 1 and the padded horizon") + identities = tuple(self.identities) + if len(identities) != count or not all( + isinstance(value, CandidateIdentity) for value in identities + ): + raise ValueError("identities must contain one CandidateIdentity per row") + keys = [(value.candidate_id, value.attempt_id) for value in identities] + if len(set(keys)) != len(keys): + raise ValueError("candidate attempts must be unique within a batch") + phases = tuple(self.phases) if self.phases else ((),) * count + if len(phases) != count: + raise ValueError("phases must contain one entry per candidate") + normalized_phases = [] + for row, length in enumerate(lengths.tolist()): + if dt[row, 0] != 0 or bool((dt[row, 1:length] <= 0).any()): + raise ValueError( + "valid dt must start at zero and use positive arrival intervals" + ) + if bool((dt[row, length:] != 0).any()): + raise ValueError("padded dt must be zero") + if not torch.equal( + positions[row, length:], + positions[row, length - 1].expand(horizon - length, joints), + ): + raise ValueError("padded positions must hold the final valid sample") + normalized_phases.append(_phases(phases[row], length)) + factors = tuple(self.factors) if self.factors else ({},) * count + if len(factors) != count: + raise ValueError("factors must contain one entry per candidate") + copied_factors = [] + for row_factors in factors: + copied = dict(row_factors) + for key, value in copied.items(): + _text(key, "factor name") + if type(value) not in (float, int, str) or ( + isinstance(value, (float, int)) and not math.isfinite(value) + ): + raise ValueError("factor values must be finite numbers or strings") + copied_factors.append(MappingProxyType(copied)) + indices = self.source_row_indices + if indices is not None: + if ( + not isinstance(indices, torch.Tensor) + or indices.shape != (count,) + or indices.dtype != torch.int64 + or bool((indices < 0).any()) + ): + raise ValueError( + "source_row_indices must be nonnegative int64 values of shape (C,)" + ) + indices = indices.detach().clone() + object.__setattr__(self, "positions", positions) + object.__setattr__(self, "dt", dt) + object.__setattr__(self, "valid_length", lengths) + object.__setattr__(self, "joint_names", names) + object.__setattr__(self, "identities", identities) + object.__setattr__(self, "phases", tuple(normalized_phases)) + object.__setattr__(self, "factors", tuple(copied_factors)) + object.__setattr__(self, "source_row_indices", indices) + + @property + def valid_mask(self) -> torch.Tensor: + """Return a fresh boolean mask selecting actual samples in each row.""" + return ( + torch.arange(self.positions.shape[1], device=self.positions.device)[None, :] + < self.valid_length[:, None] + ) + + def row(self, index: int) -> CandidateTrajectoryBatch: + """Copy one candidate and its aligned metadata, retaining padded shape.""" + _integer(index, "row index") + if index >= len(self.identities): + raise IndexError("candidate row is out of range") + return CandidateTrajectoryBatch( + positions=self.positions[index : index + 1], + dt=self.dt[index : index + 1], + valid_length=self.valid_length[index : index + 1], + identities=(self.identities[index],), + joint_names=self.joint_names, + phases=(self.phases[index],), + factors=(self.factors[index],), + source_row_indices=( + None + if self.source_row_indices is None + else self.source_row_indices[index : index + 1] + ), + ) + + +@dataclass(frozen=True) +class ValidationCheck: + """One mandatory check; unavailable and not-run checks cannot accept data.""" + + check_id: str + status: Literal["not_run", "passed", "failed", "unavailable"] + detail: str = "" + metrics: Mapping[str, float] = field(default_factory=dict) + + def __post_init__(self) -> None: + _text(self.check_id, "check_id") + if self.status not in ("not_run", "passed", "failed", "unavailable"): + raise ValueError("unknown validation status") + if not isinstance(self.detail, str): + raise ValueError("validation detail must be text") + metrics = dict(self.metrics) + for key, value in metrics.items(): + _text(key, "metric key") + if type(value) not in (int, float) or not math.isfinite(value): + raise ValueError("validation metrics must be finite numbers") + object.__setattr__(self, "metrics", MappingProxyType(metrics)) + + +@dataclass(frozen=True) +class ValidationResult: + """A fail-closed collection of required validation checks.""" + + checks: tuple[ValidationCheck, ...] + + def __post_init__(self) -> None: + checks = tuple(self.checks) + if not all(isinstance(check, ValidationCheck) for check in checks): + raise ValueError("checks must contain ValidationCheck values") + if len({check.check_id for check in checks}) != len(checks): + raise ValueError("validation check IDs must be unique") + object.__setattr__(self, "checks", checks) + + @property + def accepted(self) -> bool: + """Return whether every required check was run and passed.""" + return bool(self.checks) and all( + check.status == "passed" for check in self.checks + ) + + +@dataclass(frozen=True) +class ExpertEpisode: + """Frozen rollout evidence with causal ``obs[t], action[t], obs[t+1]`` pairs. + + Observations and timestamps contain ``T+1`` samples, including the terminal + observation. Actions contain the ``T`` commands actually submitted. This + contract checks structure; hosts remain responsible for measuring evidence. + """ + + identity: CandidateIdentity + observations: Mapping[str, torch.Tensor] + actions: torch.Tensor + timestamps: torch.Tensor + action_representation: str + validation: ValidationResult + episode_id: str + commit_id: str + phases: tuple[TrajectoryPhase, ...] = () + metadata: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.identity, CandidateIdentity) or not isinstance( + self.validation, ValidationResult + ): + raise ValueError( + "episode identity and validation must use their value contracts" + ) + _text(self.episode_id, "episode_id") + _text(self.commit_id, "commit_id") + _text(self.action_representation, "action_representation") + actions = _tensor(self.actions, "actions", 2) + times = _tensor(self.timestamps, "timestamps", 1) + steps = actions.shape[0] + if steps < 1 or actions.shape[1] < 1 or times.shape != (steps + 1,): + raise ValueError("episodes require T >= 1 actions and T+1 timestamps") + if times[0] < 0 or bool((times[1:] <= times[:-1]).any()): + raise ValueError("timestamps must be nonnegative and strictly increasing") + if not self.observations: + raise ValueError("episodes must contain actual observations") + observations = {} + for key, value in self.observations.items(): + _text(key, "observation key") + if ( + not isinstance(value, torch.Tensor) + or value.ndim < 1 + or value.shape[0] != steps + 1 + ): + raise ValueError( + "each observation must contain T+1 samples including terminal" + ) + if value.is_complex() or not bool(torch.isfinite(value).all()): + raise ValueError("observations must contain finite real values") + observations[key] = value.detach().clone() + object.__setattr__(self, "actions", actions) + object.__setattr__(self, "timestamps", times) + object.__setattr__(self, "observations", MappingProxyType(observations)) + object.__setattr__(self, "phases", _phases(self.phases, steps + 1)) + if not isinstance(self.metadata, Mapping): + raise ValueError("metadata must be a mapping") + object.__setattr__(self, "metadata", _json(self.metadata)) + + +@dataclass(frozen=True) +class CommitReceipt: + """Sink confirmation tied to stable episode/candidate IDs, never a slot. + + The idempotent ``commit_id`` remains fixed across persistence retries. + ``submission_id`` identifies an individual write attempt so a delayed + failure cannot invalidate a newer retry of the same episode. + """ + + episode_id: str + candidate_id: str + attempt_id: int + storage_id: str + commit_id: str + scene_case_id: str + confirmed: bool = True + error: str = "" + submission_id: int = 0 + + def __post_init__(self) -> None: + for name in ( + "episode_id", + "candidate_id", + "storage_id", + "commit_id", + "scene_case_id", + ): + _text(getattr(self, name), name) + _integer(self.attempt_id, "attempt_id") + _integer(self.submission_id, "submission_id") + if type(self.confirmed) is not bool: + raise ValueError("confirmed must be a boolean") + if not isinstance(self.error, str) or (self.confirmed and self.error): + raise ValueError("only unconfirmed receipts may contain an error") diff --git a/embodichain/lab/sim/motion/trajectory_augmentation/coverage.py b/embodichain/lab/sim/motion/trajectory_augmentation/coverage.py new file mode 100644 index 000000000..af0c7f1fc --- /dev/null +++ b/embodichain/lab/sim/motion/trajectory_augmentation/coverage.py @@ -0,0 +1,286 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Phase-aligned joint geometry and measured timing coverage for one scene case.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from .contracts import TrajectoryPhase + +__all__ = ["TrajectoryDescriptor", "describe_trajectory", "CoverageIndex"] + + +@dataclass(frozen=True) +class TrajectoryDescriptor: + """Measured joint geometry and elapsed times at uniform phase progress. + + Args: + phase_ids: Ordered phase identities, including contact/free mode. + geometry: Normalized joint samples, with shape ``(P, S, D)``. + timing: Elapsed seconds within each phase, with shape ``(P, S)``. + """ + + phase_ids: tuple[tuple[str, str], ...] + geometry: torch.Tensor + timing: torch.Tensor + + def __post_init__(self) -> None: + phase_ids = [] + for pair in self.phase_ids: + if ( + not isinstance(pair, (list, tuple)) + or len(pair) != 2 + or any(not isinstance(value, str) or not value for value in pair) + ): + raise ValueError("phase_ids must contain pairs of nonempty labels.") + phase_ids.append(tuple(pair)) + object.__setattr__(self, "phase_ids", tuple(phase_ids)) + geometry = self.geometry.detach().to(device="cpu", dtype=torch.float64).clone() + timing = self.timing.detach().to(device="cpu", dtype=torch.float64).clone() + if geometry.ndim != 3 or timing.shape != geometry.shape[:2]: + raise ValueError("Descriptor geometry/timing shapes must be (P,S,D)/(P,S).") + if len(self.phase_ids) != geometry.shape[0] or min(geometry.shape) < 1: + raise ValueError("Descriptor phases must match nonempty geometry.") + if not torch.isfinite(geometry).all() or not torch.isfinite(timing).all(): + raise ValueError("Descriptor values must be finite.") + object.__setattr__(self, "geometry", geometry) + object.__setattr__(self, "timing", timing) + + +def describe_trajectory( + positions: torch.Tensor, + timestamps: torch.Tensor, + joint_limits: torch.Tensor, + *, + phases: tuple[TrajectoryPhase, ...] = (), + samples_per_phase: int = 32, +) -> TrajectoryDescriptor: + """Describe an actual rollout without rewarding additional time samples. + + Geometry is sampled by normalized joint arc length within each phase. Timing + includes the phase duration, including stationary waits. Unlabelled gaps are + included as separate phases so no measured movement is silently discarded. + + Args: + positions: Actual joint positions with shape ``(T, D)``. + timestamps: Strictly increasing observation timestamps, shape ``(T,)``. + joint_limits: Finite lower/upper joint limits, shape ``(D, 2)``. + phases: Optional half-open ranges in the actual observation sequence. + samples_per_phase: Number of progress samples, at least two. + + Returns: + Owned CPU descriptor separating geometry from timing. + + Raises: + ValueError: If samples, limits, times or phase ranges are invalid. + """ + q = positions.detach().to(device="cpu", dtype=torch.float64) + t = timestamps.detach().to(device="cpu", dtype=torch.float64) + limits = joint_limits.detach().to(device="cpu", dtype=torch.float64) + if q.ndim != 2 or q.shape[0] < 2 or q.shape[1] < 1: + raise ValueError("positions must contain at least two observations.") + if t.shape != (q.shape[0],) or limits.shape != (q.shape[1], 2): + raise ValueError("timestamps or joint_limits shape does not match positions.") + if not all(torch.isfinite(x).all() for x in (q, t, limits)): + raise ValueError("Trajectory and joint limits must be finite.") + if not (t[1:] > t[:-1]).all() or not (limits[:, 1] > limits[:, 0]).all(): + raise ValueError("Times and joint-limit intervals must increase strictly.") + if type(samples_per_phase) is not int or samples_per_phase < 2: + raise ValueError("samples_per_phase must be at least two.") + q = (q - limits[:, 0]) / (limits[:, 1] - limits[:, 0]) + ranges: list[tuple[str, str, int, int]] = [] + cursor = 0 + for phase in phases: + if phase.start_index < cursor or phase.stop_index > q.shape[0]: + raise ValueError("Actual phase ranges must be ordered and in bounds.") + if phase.start_index > cursor: + ranges.append( + (f"gap_{len(ranges)}", "unlabelled", cursor, phase.start_index) + ) + ranges.append((phase.phase_id, phase.kind, phase.start_index, phase.stop_index)) + cursor = phase.stop_index + if cursor < q.shape[0]: + ranges.append((f"gap_{len(ranges)}", "unlabelled", cursor, q.shape[0])) + + geometry, timing = [], [] + for _, _, start, stop in ranges: + # Include the entering edge of a phase; otherwise disjoint sample ranges + # would leave their connecting motion out of the descriptor. + first = max(0, start - 1) + values, times = q[first:stop], t[first:stop] - t[first] + distances = torch.linalg.vector_norm(values[1:] - values[:-1], dim=-1) + arc = torch.cat((torch.zeros(1, dtype=q.dtype), distances.cumsum(0))) + if float(arc[-1]) <= 1e-12: + geometry.append(values[:1].expand(samples_per_phase, -1).clone()) + timing.append( + torch.linspace(0, float(times[-1]), samples_per_phase, dtype=q.dtype) + ) + continue + # Remove repeated arc coordinates for interpolation. The final sample's + # time is kept separately to retain terminal waits. + keep = torch.cat((torch.ones(1, dtype=torch.bool), distances > 1e-12)) + unique_arc, unique_q, unique_t = arc[keep], values[keep], times[keep] + grid = torch.linspace(0, float(arc[-1]), samples_per_phase, dtype=q.dtype) + upper = torch.searchsorted(unique_arc, grid).clamp(1, len(unique_arc) - 1) + lower = upper - 1 + weight = (grid - unique_arc[lower]) / (unique_arc[upper] - unique_arc[lower]) + geometry.append(torch.lerp(unique_q[lower], unique_q[upper], weight[:, None])) + elapsed = torch.lerp(unique_t[lower], unique_t[upper], weight) + elapsed[-1] = times[-1] + timing.append(elapsed) + return TrajectoryDescriptor( + tuple((name, kind) for name, kind, _, _ in ranges), + torch.stack(geometry), + torch.stack(timing), + ) + + +@dataclass +class _Reservation: + family_id: str + descriptor: TrajectoryDescriptor + committed: bool = False + + +class CoverageIndex: + """Reserve near-distinct measured trajectories and count confirmed geometry. + + One instance belongs to one scene case. Pending writes occupy capacity but + do not contribute to formal coverage. A timing variant cannot increase the + geometry count even if small physical tracking differences are measured. + + Args: + geometry_tolerance: Maximum normalized joint difference for near copies. + target_per_geometry: Maximum distinct timing variants per geometry. + timing_tolerance_s: Maximum timing difference for duplicate episodes. + """ + + def __init__( + self, + *, + geometry_tolerance: float = 0.01, + target_per_geometry: int = 4, + timing_tolerance_s: float = 1e-6, + ) -> None: + import math + + if any( + not math.isfinite(x) or x < 0 + for x in (geometry_tolerance, timing_tolerance_s) + ): + raise ValueError("Coverage tolerances must be finite and non-negative.") + if ( + isinstance(target_per_geometry, bool) + or not isinstance(target_per_geometry, int) + or target_per_geometry < 1 + ): + raise ValueError("target_per_geometry must be a positive integer.") + self._geometry_tolerance = geometry_tolerance + self._timing_tolerance = timing_tolerance_s + self._target = target_per_geometry + self._entries: dict[str, _Reservation] = {} + self._aliases: dict[str, str] = {} + + @property + def geometry_count(self) -> int: + """Number of distinct geometries with a confirmed persistence receipt.""" + return len( + {entry.family_id for entry in self._entries.values() if entry.committed} + ) + + def reserve( + self, + commit_id: str, + geometry_family_id: str, + descriptor: TrajectoryDescriptor, + ) -> bool: + """Reserve a distinct measured rollout before submitting its payload. + + Args: + commit_id: Stable, unique persistence identity. + geometry_family_id: Family shared by proposed timing variants. + descriptor: Descriptor calculated from the actual execution. + + Returns: + False for a near duplicate or a geometry whose quota is full. + """ + if not commit_id or not geometry_family_id: + raise ValueError("Commit and geometry-family IDs must be nonempty.") + if commit_id in self._entries: + raise ValueError("A commit ID already has a coverage reservation.") + family = self._aliases.get(geometry_family_id, geometry_family_id) + if geometry_family_id not in self._aliases: + for entry in self._entries.values(): + old = entry.descriptor + compatible = ( + old.phase_ids == descriptor.phase_ids + and old.geometry.shape == descriptor.geometry.shape + ) + if compatible and torch.allclose( + old.geometry, + descriptor.geometry, + atol=self._geometry_tolerance, + rtol=0, + ): + family = entry.family_id + break + matching = [ + entry for entry in self._entries.values() if entry.family_id == family + ] + for entry in matching: + old = entry.descriptor + compatible = ( + old.phase_ids == descriptor.phase_ids + and old.geometry.shape == descriptor.geometry.shape + ) + if compatible and torch.allclose( + old.timing, descriptor.timing, atol=self._timing_tolerance, rtol=0 + ): + return False + if len(matching) >= self._target: + return False + owned = TrajectoryDescriptor( + descriptor.phase_ids, descriptor.geometry, descriptor.timing + ) + self._entries[commit_id] = _Reservation(family, owned) + self._aliases[geometry_family_id] = family + return True + + def confirm(self, commit_id: str) -> None: + """Confirm a reserved commit; repeated confirmations are idempotent. + + Args: + commit_id: Previously reserved persistence identity. + """ + self._entries[commit_id].committed = True + + def release(self, commit_id: str) -> None: + """Release an uncommitted reservation after write failure. + + Args: + commit_id: Previously reserved persistence identity. + + Raises: + ValueError: If persistence has already been confirmed. + """ + entry = self._entries[commit_id] + if entry.committed: + raise ValueError("Committed coverage cannot be released.") + del self._entries[commit_id] diff --git a/embodichain/lab/sim/motion/trajectory_augmentation/operators.py b/embodichain/lab/sim/motion/trajectory_augmentation/operators.py new file mode 100644 index 000000000..41edc4343 --- /dev/null +++ b/embodichain/lab/sim/motion/trajectory_augmentation/operators.py @@ -0,0 +1,319 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Grasp-pose and constrained qpos operators requiring downstream validation.""" + +from __future__ import annotations + +from dataclasses import replace +import math + +import torch + +from embodichain.utils.math import axis_angle_to_rotation_matrix, pose_inv + +from .contracts import ( + TrajectoryTemplate, + ValidationCheck, + ValidationResult, + _pose, + _tensor, +) + +__all__ = [ + "rotate_grasp_about_object_axis", + "joint_residual", + "retime", + "validate_motion_limits", +] + + +def rotate_grasp_about_object_axis( + object_pose: torch.Tensor, + grasp_pose: torch.Tensor, + *, + axis: torch.Tensor, + angles: torch.Tensor, +) -> torch.Tensor: + """Rotate a reference grasp about an axis through a fixed object's origin. + + The axis is expressed in the object's local frame. Both the grasp position + and orientation rotate about that axis; the object itself is unchanged. + For example, a cube permits quarter-turn variants of a centered top grasp. + The caller owns object symmetry and grasp validity: arbitrary rotations need + not preserve contacts. Every candidate requires fresh IK/path planning and + physical grasp validation. This function does not modify a qpos template. + + Args: + object_pose: Object transform in the scene frame, shape ``(4, 4)``. + grasp_pose: Reference TCP transform in that same frame, shape ``(4, 4)``. + axis: Finite nonzero object-local rotation axis, shape ``(3,)``. + angles: Finite rotation angles in radians, shape ``(C,)``, with C > 0. + + Returns: + Owned candidate TCP transforms with shape ``(C, 4, 4)``, using the + object pose's device and dtype, in the original scene frame. + + Raises: + ValueError: If transforms, axis, or angles are malformed or nonfinite. + """ + source_object = _pose(object_pose, "object_pose") + source_grasp = _pose(grasp_pose, "grasp_pose") + direction = _tensor(axis, "axis", 1).to(source_object.device, torch.float64) + rotations = _tensor(angles, "angles", 1).to(source_object.device, torch.float64) + if direction.shape != (3,) or rotations.numel() == 0: + raise ValueError("axis must have shape (3,) and angles must be nonempty") + norm = torch.linalg.vector_norm(direction) + if not torch.isfinite(norm) or norm <= 0: + raise ValueError("axis must have a finite nonzero norm") + # Calculate in float64 to avoid overflow in low-precision axis norms and + # accumulated transforms. Wrapping also bounds the Rodrigues input. + rotations = torch.remainder(rotations + math.pi, 2 * math.pi) - math.pi + transforms = torch.eye(4, device=source_object.device, dtype=torch.float64).repeat( + rotations.numel(), 1, 1 + ) + transforms[:, :3, :3] = axis_angle_to_rotation_matrix( + rotations[:, None] * (direction / norm) + ) + anchor = source_object.to(dtype=torch.float64) + reference = source_grasp.to(device=anchor.device, dtype=anchor.dtype) + return (anchor @ transforms @ pose_inv(anchor) @ reference).to(source_object) + + +def _allowed_phases(template: TrajectoryTemplate, operator: str) -> list: + if operator not in template.allowed_operators: + raise ValueError(f"Template does not allow {operator}.") + phases = [ + phase + for phase in template.phases + if phase.kind == "free" and operator in phase.allowed_operators + ] + if not phases: + raise ValueError(f"No explicitly annotated free phase allows {operator}.") + if not template.controlled_joint_indices: + raise ValueError("Augmentation requires explicit controlled_joint_indices.") + return phases + + +def joint_residual( + template: TrajectoryTemplate, + *, + joint_limits: torch.Tensor, + normalized_scale: float, + generator: torch.Generator, +) -> TrajectoryTemplate: + """Add one smooth, endpoint-preserving residual per allowed free phase. + + This operator applies only to explicitly permitted joint-space motion. + Cartesian contact constraints require a planning adapter and cannot be + inferred from qpos. It changes only the declared controlled joints. + + Args: + template: Annotated full-joint reference. + joint_limits: Finite lower/upper limits with shape ``(D, 2)``. + normalized_scale: Maximum offset as a fraction of each joint range. + generator: Explicit local generator; global RNG is never consumed. + + Returns: + A new template requiring path and physical validation. + + Raises: + ValueError: If permissions, limits or resulting positions are invalid. + """ + phases = _allowed_phases(template, "joint_residual") + if not math.isfinite(normalized_scale) or not 0 <= normalized_scale <= 1: + raise ValueError("normalized_scale must be finite and within [0, 1].") + q = template.positions.clone() + limits = joint_limits.to(device=q.device, dtype=q.dtype) + if ( + limits.shape != (q.shape[1], 2) + or not torch.isfinite(limits).all() + or not (limits[:, 1] > limits[:, 0]).all() + ): + raise ValueError("joint_limits must be finite increasing intervals (D,2).") + indices = list(template.controlled_joint_indices) + for phase in phases: + length = phase.stop_index - phase.start_index + if length < 3: + raise ValueError("A spatially variable phase needs an interior sample.") + progress = torch.linspace(0, 1, length, device=q.device, dtype=q.dtype) + # This basis has zero value and derivative at both phase endpoints. + envelope = 16 * progress.square() * (1 - progress).square() + values = torch.rand( + len(indices), generator=generator, device=generator.device, dtype=q.dtype + ).to(q.device) + offsets = ( + (2 * values - 1) + * normalized_scale + * (limits[indices, 1] - limits[indices, 0]) + ) + q[phase.start_index : phase.stop_index, indices] += envelope[:, None] * offsets + if ((q < limits[:, 0]) | (q > limits[:, 1])).any(): + raise ValueError( + "Sampled residual violates joint limits; reject this proposal." + ) + return replace(template, positions=q) + + +def retime( + template: TrajectoryTemplate, + *, + duration_scale: float, + control_dt: float, + max_samples: int = 1_000_000, +) -> TrajectoryTemplate: + """Retime permitted free phases and resample onto the host control clock. + + Contact and wait durations, phase endpoints, and inter-phase intervals are + preserved. Phase boundaries must already align with the host clock. Scaled + phase durations round up to whole control steps; sample counts therefore + change. The function does not claim dynamic or collision validity. + + Args: + template: Reference containing explicitly permitted ``retime`` phases. + duration_scale: Positive factor; greater than one slows free motion. + control_dt: Authoritative host control period in seconds. + max_samples: Allocation bound for the resulting trajectory. + + Returns: + A new uniformly timed template with remapped phase indices. + + Raises: + ValueError: If permissions, timing or allocation constraints fail. + """ + phases = _allowed_phases(template, "retime") + if not all(math.isfinite(x) and x > 0 for x in (duration_scale, control_dt)): + raise ValueError( + "Duration scale and control period must be positive and finite." + ) + if type(max_samples) is not int or max_samples < 2: + raise ValueError("max_samples must be an integer of at least two.") + q = template.positions + intervals = template.dt.to(dtype=torch.float64).clone() + original_times = intervals.cumsum(0) + anchors = {0, q.shape[0] - 1} + variable_interior = set() + for phase in phases: + variable_interior.update(range(phase.start_index + 1, phase.stop_index - 1)) + # Every contact, wait and unlabelled sample is a locked event. Refuse a + # coarser clock that would silently remove an intermediate tool command. + anchors.update(i for i in range(q.shape[0]) if i not in variable_interior) + for phase in template.phases: + anchors.update((phase.start_index, phase.stop_index - 1)) + for index in anchors: + steps = float(original_times[index]) / control_dt + if not math.isclose(steps, round(steps), abs_tol=1e-5, rel_tol=1e-6): + raise ValueError("Phase anchors must align with the host control clock.") + uncontrolled = [ + i for i in range(q.shape[1]) if i not in template.controlled_joint_indices + ] + for phase in phases: + start, stop = phase.start_index, phase.stop_index + if stop - start < 2: + raise ValueError("A retimed phase needs at least two samples.") + if uncontrolled and not torch.equal( + q[start:stop, uncontrolled], + q[start : start + 1, uncontrolled].expand(stop - start, -1), + ): + raise ValueError( + "Retiming may not change an uncontrolled joint's schedule." + ) + duration = float(intervals[start + 1 : stop].sum()) + steps = max(1, math.ceil(duration * duration_scale / control_dt - 1e-6)) + intervals[start + 1 : stop] *= steps * control_dt / duration + times = intervals.cumsum(0) + count = round(float(times[-1]) / control_dt) + 1 + if count < 2 or count > max_samples: + raise ValueError("Retimed trajectory exceeds the declared sample budget.") + grid = torch.arange(count, dtype=times.dtype, device=times.device) * control_dt + upper = torch.searchsorted(times, grid).clamp(1, q.shape[0] - 1) + lower = upper - 1 + weight = ((grid - times[lower]) / (times[upper] - times[lower])).to(q.dtype) + result = torch.lerp(q[lower], q[upper], weight[:, None]) + # Preserve the exact declared anchors, including closed-contact endpoints. + for index in anchors: + result[round(float(times[index]) / control_dt)] = q[index] + remapped = tuple( + replace( + phase, + start_index=round(float(times[phase.start_index]) / control_dt), + stop_index=round(float(times[phase.stop_index - 1]) / control_dt) + 1, + ) + for phase in template.phases + ) + dt = torch.full((count,), control_dt, dtype=template.dt.dtype, device=q.device) + dt[0] = 0 + return replace(template, positions=result, dt=dt, phases=remapped) + + +def validate_motion_limits( + template: TrajectoryTemplate, + *, + velocity_limits: torch.Tensor, + acceleration_limits: torch.Tensor, +) -> ValidationResult: + """Check sampled finite-difference motion limits without asserting task success. + + Args: + template: Timed full-joint path to check. + velocity_limits: Positive per-joint speed bounds, shape ``(D,)``. + acceleration_limits: Positive per-joint acceleration bounds, shape ``(D,)``. + + Returns: + A single ``motion_limits`` check with maximum normalized ratios. + """ + # Accumulate derivatives outside the source's storage dtype. Half-precision + # ratios can overflow even for ordinary joint motion and finite profiles. + q = template.positions.to(dtype=torch.float64) + dt = template.dt.to(dtype=torch.float64) + limits = [ + x.to(device=q.device, dtype=q.dtype) + for x in (velocity_limits, acceleration_limits) + ] + if any( + x.shape != (q.shape[1],) or not torch.isfinite(x).all() or not (x > 0).all() + for x in limits + ): + raise ValueError("Dynamic limits must be positive finite per-joint vectors.") + velocity = (q[1:] - q[:-1]) / dt[1:, None] + acceleration = (velocity[1:] - velocity[:-1]) / ((dt[2:] + dt[1:-1])[:, None] / 2) + speed_ratio = float((velocity.abs() / limits[0]).max()) if velocity.numel() else 0.0 + acceleration_ratio = ( + float((acceleration.abs() / limits[1]).max()) if acceleration.numel() else 0.0 + ) + passed = ( + math.isfinite(speed_ratio) + and math.isfinite(acceleration_ratio) + and max(speed_ratio, acceleration_ratio) <= 1 + 1e-6 + ) + metrics = { + name: value + for name, value in ( + ("speed_ratio", speed_ratio), + ("acceleration_ratio", acceleration_ratio), + ) + if math.isfinite(value) + } + return ValidationResult( + ( + ValidationCheck( + "motion_limits", + "passed" if passed else "failed", + detail="" if len(metrics) == 2 else "Non-finite sampled derivatives.", + metrics=metrics, + ), + ) + ) diff --git a/embodichain/lab/sim/motion/trajectory_augmentation/session.py b/embodichain/lab/sim/motion/trajectory_augmentation/session.py new file mode 100644 index 000000000..fa8feb77a --- /dev/null +++ b/embodichain/lab/sim/motion/trajectory_augmentation/session.py @@ -0,0 +1,637 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Host-independent candidate bookkeeping and confirmed collection quotas.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass, fields, replace +import hashlib +import json +import time +from types import MappingProxyType + +import torch + +from .cfg import TrajectoryGenerationJobCfg +from .contracts import ( + CandidateIdentity, + CandidateTrajectoryBatch, + CommitReceipt, + ExpertEpisode, + SceneCase, + ValidationCheck, + ValidationResult, +) +from .coverage import CoverageIndex, describe_trajectory + +__all__ = ["GenerationSession"] + + +@dataclass +class _Attempt: + identity: CandidateIdentity + state: str = "proposed" + trajectory: CandidateTrajectoryBatch | None = None + joint_names: tuple[str, ...] = () + episode: ExpertEpisode | None = None + episode_byte_budget: int = 0 + submission_id: int = 0 + reason: str | None = None + planning_validation: ValidationResult | None = None + rollout_validation: ValidationResult | None = None + + +def _digest(*values: object) -> str: + return hashlib.sha256( + json.dumps(values, separators=(",", ":")).encode() + ).hexdigest() + + +def _tensor_bytes(*values: torch.Tensor | None) -> int: + return sum( + value.numel() * value.element_size() for value in values if value is not None + ) + + +def _json_bytes(value: object) -> int: + def encode(item: object) -> dict: + if isinstance(item, Mapping): + return dict(item) + return {field.name: getattr(item, field.name) for field in fields(item)} + + return len(json.dumps(value, default=encode, ensure_ascii=False).encode()) + + +def _candidate_bytes(batch: CandidateTrajectoryBatch) -> int: + return _tensor_bytes( + batch.positions, batch.dt, batch.valid_length, batch.source_row_indices + ) + sum( + _json_bytes((identity, batch.joint_names, phases, factors)) + for identity, phases, factors in zip( + batch.identities, batch.phases, batch.factors + ) + ) + + +def _validation_summary(result: ValidationResult) -> ValidationResult: + """Keep four-state outcomes with bounded text and numeric audit context.""" + if len(result.checks) > 64: + raise ValueError("Validation supports at most 64 checks per result.") + return ValidationResult( + tuple( + ValidationCheck( + check.check_id, + check.status, + check.detail[:256], + dict(tuple(check.metrics.items())[:16]), + ) + for check in result.checks + ) + ) + + +class GenerationSession: + """Own one job's bounded candidates, budgets, evidence, and commit history. + + This value-only session never reads an environment or steps a simulator. + Call ``mark_rollout_started`` only after the first actual command. A sink + must provide one final outcome per submission and persist ``commit_id`` + idempotently across retries. Contradictory final receipts are rejected. + The caller serializes operations. Cancellation means stopping new work and + releasing individual candidates; pending writes still require final receipts. + Byte limits include tensor data and serialized value metadata, excluding + Python object allocation overhead. The runner separately bounds host buffers + and any copies it allocates outside this session. + + Args: + cfg: Semantically valid job settings, copied on construction. + clock: Monotonic seconds, injectable for deterministic budget tests. + """ + + def __init__( + self, + cfg: TrajectoryGenerationJobCfg, + *, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self._cfg = TrajectoryGenerationJobCfg.from_mapping(cfg.to_dict()) + self._clock, self._started_at = clock, clock() + self._cases: dict[ + tuple[str, str], tuple[SceneCase, torch.Tensor, tuple[str, ...]] + ] = {} + self._coverage: dict[str, CoverageIndex] = {} + self._ordinals: dict[tuple[str, ...], int] = {} + self._attempts: dict[str, _Attempt] = {} + self._commits: dict[str, str] = {} + self._receipts: dict[tuple[str, int], CommitReceipt] = {} + self._counts = dict.fromkeys( + ( + "proposed", + "planned_valid", + "rollout_attempted", + "validated_accepted", + "committed", + ), + 0, + ) + + def register_case( + self, + case: SceneCase, + joint_limits: torch.Tensor, + *, + joint_names: tuple[str, ...], + ) -> None: + """Register immutable case conditions and complete joint normalization limits.""" + names = tuple(joint_names) + if ( + isinstance(joint_names, str) + or not names + or any(not isinstance(name, str) or not name.strip() for name in names) + or len(set(names)) != len(names) + ): + raise ValueError("joint_names must contain unique nonempty names.") + limits = joint_limits.detach().to(device="cpu", dtype=torch.float64).clone() + if limits.shape != (len(names), 2): + raise ValueError("joint_limits must have shape (len(joint_names), 2).") + if not torch.isfinite(limits).all() or not (limits[:, 1] > limits[:, 0]).all(): + raise ValueError( + "joint_limits must be finite, strictly increasing intervals." + ) + key = (case.scene_case_id, case.initial_state_id) + if key in self._cases: + previous, previous_limits, previous_names = self._cases[key] + if ( + previous != case + or not torch.equal(previous_limits, limits) + or previous_names != names + ): + raise ValueError("A registered case's fixed conditions cannot change.") + return + for (case_id, _), ( + previous, + previous_limits, + previous_names, + ) in self._cases.items(): + if case_id == case.scene_case_id and ( + replace(previous, initial_state_id=case.initial_state_id) != case + or not torch.equal(previous_limits, limits) + or previous_names != names + ): + raise ValueError( + "Initial states in one case must share scene and robot conditions." + ) + self._cases[key] = case, limits, names + if case.scene_case_id not in self._coverage: + coverage = self._cfg.augmentation.coverage + self._coverage[case.scene_case_id] = CoverageIndex( + geometry_tolerance=coverage.joint_dedup_normalized_tol, + target_per_geometry=coverage.target_per_cell, + ) + + def propose( + self, + case_id: str, + initial_state_id: str, + *, + source_id: str, + source_revision: str, + template_id: str, + operator_id: str, + geometry_family_id: str | None = None, + parent_id: str | None = None, + ) -> tuple[CandidateIdentity, torch.Generator]: + """Allocate stable lineage and a local CPU RNG independent of execution slots.""" + self._cases[(case_id, initial_state_id)] + if ( + self.stop_reason + or self._counts["proposed"] >= self._cfg.collection.max_proposals + ): + raise RuntimeError( + "Proposal budget is exhausted or the session has stopped." + ) + if not isinstance(operator_id, str) or not operator_id.strip(): + raise ValueError("operator_id must be a nonempty string.") + stream = ( + case_id, + initial_state_id, + source_id, + source_revision, + template_id, + operator_id, + ) + ordinal = self._ordinals.get(stream, 0) + digest = _digest(self._cfg.augmentation.seed, *stream, ordinal) + identity = CandidateIdentity( + case_id, + initial_state_id, + "candidate_" + digest, + geometry_family_id or "geometry_" + digest, + source_id, + source_revision, + template_id, + parent_id=parent_id, + ) + if parent_id is not None: + parent = self._attempts[parent_id].identity + if (parent.scene_case_id, parent.initial_state_id) != ( + case_id, + initial_state_id, + ): + raise ValueError( + "A candidate parent must belong to the same case and initial state." + ) + self._ordinals[stream] = ordinal + 1 + self._attempts[identity.candidate_id] = _Attempt(identity) + self._counts["proposed"] += 1 + return identity, torch.Generator(device="cpu").manual_seed( + int(digest[:16], 16) % 2**63 + ) + + def _attempt(self, identity: CandidateIdentity, *states: str) -> _Attempt: + attempt = self._attempts[identity.candidate_id] + if attempt.identity != identity or (states and attempt.state not in states): + raise ValueError( + "Candidate identity or lifecycle state does not match this operation." + ) + return attempt + + def _validation_passes( + self, validation: ValidationResult, *, rollout: bool + ) -> bool: + required = ( + {"path_collision"} if self._cfg.validation.require_path_collision else set() + ) + if rollout and self._cfg.validation.require_task_success: + required.add("task_success") + return validation.accepted and required <= { + check.check_id for check in validation.checks + } + + def add_planned( + self, batch: CandidateTrajectoryBatch, validation: ValidationResult + ) -> None: + """Atomically admit validated candidate rows within ready count and byte limits. + + The caller owns physical planning checks. Nonempty, passing checks must + include ``path_collision`` when required by this job. + """ + attempts = [ + self._attempt(identity, "proposed") for identity in batch.identities + ] + for attempt in attempts: + attempt.planning_validation = _validation_summary(validation) + if not self._validation_passes(validation, rollout=False): + for attempt in attempts: + attempt.reason = "planning_validation_failed" + raise ValueError( + "Planned candidates require passing mandatory planning checks." + ) + for identity in batch.identities: + names = self._cases[(identity.scene_case_id, identity.initial_state_id)][2] + if batch.joint_names != names: + raise ValueError( + "Candidate joint layout does not match the registered limits." + ) + ready = [ + attempt.trajectory + for attempt in self._attempts.values() + if attempt.state == "ready" + ] + if ( + len(ready) + len(batch.identities) + > self._cfg.execution.ready_high_watermark + or sum(map(_candidate_bytes, ready)) + _candidate_bytes(batch) + > self._cfg.execution.ready_max_bytes + ): + raise BufferError("The ready candidate pool is full.") + rows = [batch.row(index) for index in range(len(batch.identities))] + for row in rows: + attempt = self._attempt(row.identities[0]) + attempt.reason = None + attempt.state, attempt.trajectory, attempt.joint_names = ( + "ready", + row, + row.joint_names, + ) + self._counts["planned_valid"] += len(rows) + + def _reserved(self) -> int: + return sum( + attempt.state in {"assigned", "running", "pending_write", "committed"} + for attempt in self._attempts.values() + ) + + def take_ready( + self, + case_id: str, + initial_state_id: str, + *, + episode_byte_budget: int | None = None, + ) -> CandidateTrajectoryBatch | None: + """Reserve a row, rollout attempt, and pending payload capacity before execution. + + ``episode_byte_budget`` is a trusted upper bound on serialized metadata + plus tensor data in the eventual episode. ``None`` reserves the entire + pending byte limit, conservatively allowing only one outstanding episode. + Smaller explicit bounds permit concurrent rows; exceeding a bound rejects + the episode. The runner owns separate limits on host-side allocations. + """ + self._cases[(case_id, initial_state_id)] + byte_limit = self._cfg.persistence.pending_max_bytes + budget = byte_limit if episode_byte_budget is None else episode_byte_budget + if type(budget) is not int or not 0 < budget <= byte_limit: + raise ValueError( + "episode_byte_budget must be a positive integer within pending_max_bytes." + ) + pending_count, pending_bytes = self._pending_reservation() + assigned = sum( + attempt.state == "assigned" for attempt in self._attempts.values() + ) + if ( + self.stop_reason + or self._reserved() >= self._cfg.collection.target_committed_episodes + or self._counts["rollout_attempted"] + assigned + >= self._cfg.collection.max_rollout_attempts + or pending_count >= self._cfg.persistence.pending_episode_limit + or pending_bytes + budget > byte_limit + ): + return None + for attempt in self._attempts.values(): + identity = attempt.identity + if attempt.state == "ready" and ( + identity.scene_case_id, + identity.initial_state_id, + ) == (case_id, initial_state_id): + attempt.state = "assigned" + attempt.episode_byte_budget = budget + trajectory, attempt.trajectory = attempt.trajectory, None + return trajectory + return None + + def mark_rollout_started(self, identity: CandidateIdentity) -> None: + """Count an assigned candidate once its first command actually ran.""" + attempt = self._attempt(identity, "assigned") + attempt.state, attempt.trajectory = "running", None + self._counts["rollout_attempted"] += 1 + + def release( + self, identity: CandidateIdentity, *, reason: str = "discarded" + ) -> None: + """Discard uncommitted work while preserving its failure reason and checks.""" + if not isinstance(reason, str) or not reason.strip(): + raise ValueError("release reason must be a nonempty string.") + attempt = self._attempt( + identity, "proposed", "ready", "assigned", "running", "write_failed" + ) + attempt.state, attempt.trajectory, attempt.episode = "released", None, None + attempt.episode_byte_budget = 0 + attempt.reason = reason + + def episode_ids(self, identity: CandidateIdentity) -> tuple[str, str]: + """Return the deterministic episode and idempotent commit IDs for an attempt.""" + self._attempt(identity) + digest = _digest(identity.candidate_id, identity.attempt_id) + return "episode_" + digest, "commit_" + digest + + def _episode_bytes(self, episode: ExpertEpisode) -> int: + return _tensor_bytes( + episode.actions, episode.timestamps, *episode.observations.values() + ) + _json_bytes( + ( + episode.identity, + episode.episode_id, + episode.commit_id, + episode.action_representation, + tuple(episode.observations), + episode.validation, + episode.phases, + episode.metadata, + ) + ) + + def _pending_reservation(self) -> tuple[int, int]: + count, size = 0, 0 + for attempt in self._attempts.values(): + if attempt.episode is not None: + count += 1 + size += self._episode_bytes(attempt.episode) + elif attempt.state in {"assigned", "running"}: + count += 1 + size += attempt.episode_byte_budget + return count, size + + def _reserve_coverage(self, episode: ExpertEpisode) -> bool: + identity = episode.identity + limits = self._cases[(identity.scene_case_id, identity.initial_state_id)][1] + descriptor = describe_trajectory( + episode.observations["joint_positions"], + episode.timestamps, + limits, + phases=episode.phases, + samples_per_phase=self._cfg.augmentation.coverage.geometry_samples_per_phase, + ) + return self._coverage[identity.scene_case_id].reserve( + episode.commit_id, identity.geometry_family_id, descriptor + ) + + def accept_episode(self, episode: ExpertEpisode) -> bool: + """Reserve passed measured evidence for writing; rejected evidence releases its slot. + + Joint positions, observation timestamps, and actual phase intervals + supply coverage. Planned targets are never used as measured evidence. + """ + attempt = self._attempt(episode.identity, "running") + if (episode.episode_id, episode.commit_id) != self.episode_ids( + episode.identity + ): + raise ValueError("Episode and commit IDs must match the candidate attempt.") + attempt.rollout_validation = _validation_summary(episode.validation) + if episode.action_representation != "qpos" or episode.actions.shape[1] != len( + attempt.joint_names + ): + raise ValueError( + "Episode actions must use full-joint qpos labels in candidate joint order." + ) + if not self._validation_passes(episode.validation, rollout=True): + self.release(episode.identity, reason="rollout_validation_failed") + return False + positions = episode.observations.get("joint_positions") + if ( + positions is None + or positions.ndim != 2 + or positions.shape[1] != len(attempt.joint_names) + ): + raise ValueError( + "Evidence requires actual joint_positions in the candidate joint order." + ) + if self._episode_bytes(episode) > attempt.episode_byte_budget: + self.release(episode.identity, reason="episode_byte_budget_exceeded") + raise BufferError("Payload exceeds its reserved episode byte budget.") + owned = replace(episode) + if not self._reserve_coverage(owned): + self.release(episode.identity, reason="coverage_rejected") + return False + attempt.state, attempt.episode = "pending_write", owned + attempt.episode_byte_budget = 0 + self._commits[episode.commit_id] = episode.identity.candidate_id + self._counts["validated_accepted"] += 1 + return True + + def apply_receipt(self, receipt: CommitReceipt) -> bool: + """Apply a final submission result to its original identity, idempotently. + + A repeated receipt is ignored. Different outcomes for one submission + are invalid even after retry; the sink must resolve ambiguous writes + before reporting a final failure and releasing its reservation. + """ + attempt = self._attempts[self._commits[receipt.commit_id]] + identity = attempt.identity + expected = ( + *self.episode_ids(identity), + identity.candidate_id, + identity.attempt_id, + identity.scene_case_id, + ) + if ( + receipt.episode_id, + receipt.commit_id, + receipt.candidate_id, + receipt.attempt_id, + receipt.scene_case_id, + ) != expected: + raise ValueError("Receipt identity does not match its original candidate.") + key = (receipt.commit_id, receipt.submission_id) + if key in self._receipts: + if self._receipts[key] != receipt: + raise ValueError( + "Contradictory final receipts for one write submission." + ) + return False + if ( + attempt.state != "pending_write" + or receipt.submission_id != attempt.submission_id + ): + raise ValueError("Receipt does not identify a pending write submission.") + coverage = self._coverage[identity.scene_case_id] + if receipt.confirmed: + coverage.confirm(receipt.commit_id) + attempt.state, attempt.episode = "committed", None + self._counts["committed"] += 1 + else: + coverage.release(receipt.commit_id) + attempt.state = "write_failed" + attempt.reason = receipt.error or "write_failed" + self._receipts[key] = receipt + return True + + def retry_write(self, commit_id: str) -> tuple[ExpertEpisode, int]: + """Re-reserve a failed payload and return a copy plus its next submission ID. + + A failed write frees collection and coverage quota. Retry therefore + fails if another candidate has occupied either quota in the meantime. + It never adds a proposal, rollout attempt, or accepted-episode count. + """ + attempt = self._attempts[self._commits[commit_id]] + if attempt.state != "write_failed": + raise ValueError("Only a final failed write can be retried.") + if ( + self._reserved() >= self._cfg.collection.target_committed_episodes + or self._clock() - self._started_at >= self._cfg.collection.max_wall_time_s + ): + raise RuntimeError( + "No collection quota or time remains for this write retry." + ) + if not self._reserve_coverage(attempt.episode): + raise RuntimeError( + "Coverage quota is no longer available for this write retry." + ) + attempt.state = "pending_write" + attempt.reason = None + attempt.submission_id += 1 + return replace(attempt.episode), attempt.submission_id + + @property + def stop_reason(self) -> str | None: + """Report exhausted collection budgets while allowing outstanding results to finish.""" + collection = self._cfg.collection + if self._counts["committed"] >= collection.target_committed_episodes: + return "target_reached" + if self._clock() - self._started_at >= collection.max_wall_time_s: + return "wall_time_exhausted" + active = any( + attempt.state in {"assigned", "running", "pending_write"} + for attempt in self._attempts.values() + ) + if ( + not active + and self._counts["rollout_attempted"] >= collection.max_rollout_attempts + ): + return "rollout_budget_exhausted" + if ( + not active + and self._counts["proposed"] >= collection.max_proposals + and not any( + attempt.state in {"proposed", "ready"} + for attempt in self._attempts.values() + ) + ): + return "proposal_budget_exhausted" + return None + + def snapshot(self) -> Mapping[str, object]: + """Return read-only counters, case coverage, and lightweight lifecycle audit. + + Diagnostics retain each check's status, first 256 detail characters, + and first 16 numeric metrics. Results support at most 64 checks. + Pending reservations cover assigned/running upper bounds and retained + payload bytes, including failed writes that remain available for retry. + """ + counts = dict(self._counts) + pending_count, pending_bytes = self._pending_reservation() + for state in ("ready", "assigned", "running", "pending_write", "write_failed"): + counts[state] = sum( + attempt.state == state for attempt in self._attempts.values() + ) + return MappingProxyType( + { + "counts": MappingProxyType(counts), + "pending_reserved_episodes": pending_count, + "pending_reserved_bytes": pending_bytes, + "coverage": MappingProxyType( + {key: value.geometry_count for key, value in self._coverage.items()} + ), + "audit": tuple( + (attempt.identity, attempt.state, attempt.submission_id) + for attempt in self._attempts.values() + ), + "diagnostics": MappingProxyType( + { + candidate_id: MappingProxyType( + { + "reason": attempt.reason, + "planning_validation": attempt.planning_validation, + "rollout_validation": attempt.rollout_validation, + } + ) + for candidate_id, attempt in self._attempts.items() + } + ), + "stop_reason": self.stop_reason, + } + ) diff --git a/embodichain/lab/sim/workspace/__init__.py b/embodichain/lab/sim/motion/workspace/__init__.py similarity index 100% rename from embodichain/lab/sim/workspace/__init__.py rename to embodichain/lab/sim/motion/workspace/__init__.py diff --git a/embodichain/lab/sim/workspace/analyzer.py b/embodichain/lab/sim/motion/workspace/analyzer.py similarity index 99% rename from embodichain/lab/sim/workspace/analyzer.py rename to embodichain/lab/sim/motion/workspace/analyzer.py index 292aae18d..5478e7239 100644 --- a/embodichain/lab/sim/workspace/analyzer.py +++ b/embodichain/lab/sim/motion/workspace/analyzer.py @@ -38,7 +38,7 @@ from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.objects.robot import Robot -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( CacheConfig, DimensionConstraint, SamplingConfig, @@ -46,16 +46,16 @@ VisualizationConfig, MetricConfig, ) -from embodichain.lab.sim.workspace.samplers import ( +from embodichain.lab.sim.motion.workspace.samplers import ( SamplerFactory, BaseSampler, ) -from embodichain.lab.sim.workspace.caches import CacheManager -from embodichain.lab.sim.workspace.caches.results_cache import ( +from embodichain.lab.sim.motion.workspace.caches import CacheManager +from embodichain.lab.sim.motion.workspace.caches.results_cache import ( ResultsCache, compute_cache_key, ) -from embodichain.lab.sim.workspace.constraints import ( +from embodichain.lab.sim.motion.workspace.constraints import ( WorkspaceConstraintChecker, ) @@ -341,7 +341,7 @@ def _compute_dynamic_workspace_bounds(self) -> torch.Tensor: logger.log_info("Computing workspace bounds dynamically from joint space FK...") # Create a temporary sampler without constraints for initial FK computation - from embodichain.lab.sim.workspace.samplers import ( + from embodichain.lab.sim.motion.workspace.samplers import ( RandomSampler, ) @@ -1079,7 +1079,7 @@ def compute_reachability( total_reachable = 0 # Prepare random seeds for all attempts - from embodichain.lab.sim.workspace.samplers import ( + from embodichain.lab.sim.motion.workspace.samplers import ( RandomSampler, ) @@ -2004,7 +2004,7 @@ def visualize( ) # Create visualizer using factory pattern - from embodichain.lab.sim.workspace.visualizers import ( + from embodichain.lab.sim.motion.workspace.visualizers import ( VisualizerFactory, ) diff --git a/embodichain/lab/sim/workspace/caches/__init__.py b/embodichain/lab/sim/motion/workspace/caches/__init__.py similarity index 77% rename from embodichain/lab/sim/workspace/caches/__init__.py rename to embodichain/lab/sim/motion/workspace/caches/__init__.py index f8e90935d..1f3e06c2d 100644 --- a/embodichain/lab/sim/workspace/caches/__init__.py +++ b/embodichain/lab/sim/motion/workspace/caches/__init__.py @@ -19,15 +19,15 @@ Provides in-memory and disk caches (``MemoryCache``, ``DiskCache``), a ``CacheManager``, a ``ResultsCache``, and ``compute_cache_key``. """ -from embodichain.lab.sim.workspace.caches.base_cache import BaseCache -from embodichain.lab.sim.workspace.caches.memory_cache import ( +from embodichain.lab.sim.motion.workspace.caches.base_cache import BaseCache +from embodichain.lab.sim.motion.workspace.caches.memory_cache import ( MemoryCache, ) -from embodichain.lab.sim.workspace.caches.disk_cache import DiskCache -from embodichain.lab.sim.workspace.caches.cache_manager import ( +from embodichain.lab.sim.motion.workspace.caches.disk_cache import DiskCache +from embodichain.lab.sim.motion.workspace.caches.cache_manager import ( CacheManager, ) -from embodichain.lab.sim.workspace.caches.results_cache import ( +from embodichain.lab.sim.motion.workspace.caches.results_cache import ( DEFAULT_RESULTS_CACHE_DIR, ResultsCache, compute_cache_key, diff --git a/embodichain/lab/sim/workspace/caches/base_cache.py b/embodichain/lab/sim/motion/workspace/caches/base_cache.py similarity index 100% rename from embodichain/lab/sim/workspace/caches/base_cache.py rename to embodichain/lab/sim/motion/workspace/caches/base_cache.py diff --git a/embodichain/lab/sim/workspace/caches/cache_manager.py b/embodichain/lab/sim/motion/workspace/caches/cache_manager.py similarity index 91% rename from embodichain/lab/sim/workspace/caches/cache_manager.py rename to embodichain/lab/sim/motion/workspace/caches/cache_manager.py index 64428a070..f105002e7 100644 --- a/embodichain/lab/sim/workspace/caches/cache_manager.py +++ b/embodichain/lab/sim/motion/workspace/caches/cache_manager.py @@ -16,12 +16,12 @@ from typing import Literal -from embodichain.lab.sim.workspace.caches.base_cache import BaseCache -from embodichain.lab.sim.workspace.caches.memory_cache import ( +from embodichain.lab.sim.motion.workspace.caches.base_cache import BaseCache +from embodichain.lab.sim.motion.workspace.caches.memory_cache import ( MemoryCache, ) -from embodichain.lab.sim.workspace.caches.disk_cache import DiskCache -from embodichain.lab.sim.workspace.configs.cache_config import ( +from embodichain.lab.sim.motion.workspace.caches.disk_cache import DiskCache +from embodichain.lab.sim.motion.workspace.configs.cache_config import ( CacheConfig, ) diff --git a/embodichain/lab/sim/workspace/caches/cache_utils.py b/embodichain/lab/sim/motion/workspace/caches/cache_utils.py similarity index 100% rename from embodichain/lab/sim/workspace/caches/cache_utils.py rename to embodichain/lab/sim/motion/workspace/caches/cache_utils.py diff --git a/embodichain/lab/sim/workspace/caches/disk_cache.py b/embodichain/lab/sim/motion/workspace/caches/disk_cache.py similarity index 98% rename from embodichain/lab/sim/workspace/caches/disk_cache.py rename to embodichain/lab/sim/motion/workspace/caches/disk_cache.py index 7743729ef..36ce7e6d6 100644 --- a/embodichain/lab/sim/workspace/caches/disk_cache.py +++ b/embodichain/lab/sim/motion/workspace/caches/disk_cache.py @@ -21,7 +21,7 @@ from tqdm import tqdm -from embodichain.lab.sim.workspace.caches.base_cache import BaseCache +from embodichain.lab.sim.motion.workspace.caches.base_cache import BaseCache from embodichain.utils import logger diff --git a/embodichain/lab/sim/workspace/caches/memory_cache.py b/embodichain/lab/sim/motion/workspace/caches/memory_cache.py similarity index 96% rename from embodichain/lab/sim/workspace/caches/memory_cache.py rename to embodichain/lab/sim/motion/workspace/caches/memory_cache.py index 3e17fc4a5..2a2bcbd64 100644 --- a/embodichain/lab/sim/workspace/caches/memory_cache.py +++ b/embodichain/lab/sim/motion/workspace/caches/memory_cache.py @@ -18,7 +18,7 @@ from typing import List import numpy as np -from embodichain.lab.sim.workspace.caches.base_cache import BaseCache +from embodichain.lab.sim.motion.workspace.caches.base_cache import BaseCache all = [ "MemoryCache", diff --git a/embodichain/lab/sim/workspace/caches/results_cache.py b/embodichain/lab/sim/motion/workspace/caches/results_cache.py similarity index 100% rename from embodichain/lab/sim/workspace/caches/results_cache.py rename to embodichain/lab/sim/motion/workspace/caches/results_cache.py diff --git a/embodichain/lab/sim/workspace/cfg.py b/embodichain/lab/sim/motion/workspace/cfg.py similarity index 100% rename from embodichain/lab/sim/workspace/cfg.py rename to embodichain/lab/sim/motion/workspace/cfg.py diff --git a/embodichain/lab/sim/workspace/configs/__init__.py b/embodichain/lab/sim/motion/workspace/configs/__init__.py similarity index 81% rename from embodichain/lab/sim/workspace/configs/__init__.py rename to embodichain/lab/sim/motion/workspace/configs/__init__.py index fd47c0fa7..49be0b945 100644 --- a/embodichain/lab/sim/workspace/configs/__init__.py +++ b/embodichain/lab/sim/motion/workspace/configs/__init__.py @@ -19,21 +19,21 @@ Covers sampling (``SamplingConfig``), caching (``CacheConfig``), metrics (``MetricConfig`` and per-metric reachability / manipulability / density configs), visualization (``VisualizationConfig``), and dimension constraints (``DimensionConstraint``). """ -from embodichain.lab.sim.workspace.configs.cache_config import ( +from embodichain.lab.sim.motion.workspace.configs.cache_config import ( CacheConfig, ) -from embodichain.lab.sim.workspace.configs.dimension_constraint import ( +from embodichain.lab.sim.motion.workspace.configs.dimension_constraint import ( DimensionConstraint, ) -from embodichain.lab.sim.workspace.configs.sampling_config import ( +from embodichain.lab.sim.motion.workspace.configs.sampling_config import ( SamplingConfig, SamplingStrategy, ) -from embodichain.lab.sim.workspace.configs.visualization_config import ( +from embodichain.lab.sim.motion.workspace.configs.visualization_config import ( VisualizationConfig, VisualizationType, ) -from embodichain.lab.sim.workspace.configs.metric_config import ( +from embodichain.lab.sim.motion.workspace.configs.metric_config import ( MetricConfig, MetricType, ReachabilityConfig, diff --git a/embodichain/lab/sim/workspace/configs/cache_config.py b/embodichain/lab/sim/motion/workspace/configs/cache_config.py similarity index 100% rename from embodichain/lab/sim/workspace/configs/cache_config.py rename to embodichain/lab/sim/motion/workspace/configs/cache_config.py diff --git a/embodichain/lab/sim/workspace/configs/dimension_constraint.py b/embodichain/lab/sim/motion/workspace/configs/dimension_constraint.py similarity index 100% rename from embodichain/lab/sim/workspace/configs/dimension_constraint.py rename to embodichain/lab/sim/motion/workspace/configs/dimension_constraint.py diff --git a/embodichain/lab/sim/workspace/configs/metric_config.py b/embodichain/lab/sim/motion/workspace/configs/metric_config.py similarity index 100% rename from embodichain/lab/sim/workspace/configs/metric_config.py rename to embodichain/lab/sim/motion/workspace/configs/metric_config.py diff --git a/embodichain/lab/sim/workspace/configs/sampling_config.py b/embodichain/lab/sim/motion/workspace/configs/sampling_config.py similarity index 100% rename from embodichain/lab/sim/workspace/configs/sampling_config.py rename to embodichain/lab/sim/motion/workspace/configs/sampling_config.py diff --git a/embodichain/lab/sim/workspace/configs/visualization_config.py b/embodichain/lab/sim/motion/workspace/configs/visualization_config.py similarity index 100% rename from embodichain/lab/sim/workspace/configs/visualization_config.py rename to embodichain/lab/sim/motion/workspace/configs/visualization_config.py diff --git a/embodichain/lab/sim/workspace/constraints/__init__.py b/embodichain/lab/sim/motion/workspace/constraints/__init__.py similarity index 100% rename from embodichain/lab/sim/workspace/constraints/__init__.py rename to embodichain/lab/sim/motion/workspace/constraints/__init__.py diff --git a/embodichain/lab/sim/workspace/constraints/base_constraint.py b/embodichain/lab/sim/motion/workspace/constraints/base_constraint.py similarity index 100% rename from embodichain/lab/sim/workspace/constraints/base_constraint.py rename to embodichain/lab/sim/motion/workspace/constraints/base_constraint.py diff --git a/embodichain/lab/sim/workspace/constraints/workspace_constraint.py b/embodichain/lab/sim/motion/workspace/constraints/workspace_constraint.py similarity index 95% rename from embodichain/lab/sim/workspace/constraints/workspace_constraint.py rename to embodichain/lab/sim/motion/workspace/constraints/workspace_constraint.py index a19ab9594..9c599a593 100644 --- a/embodichain/lab/sim/workspace/constraints/workspace_constraint.py +++ b/embodichain/lab/sim/motion/workspace/constraints/workspace_constraint.py @@ -17,10 +17,10 @@ import numpy as np import torch -from embodichain.lab.sim.workspace.constraints.base_constraint import ( +from embodichain.lab.sim.motion.workspace.constraints.base_constraint import ( BaseConstraintChecker, ) -from embodichain.lab.sim.workspace.configs.dimension_constraint import ( +from embodichain.lab.sim.motion.workspace.configs.dimension_constraint import ( DimensionConstraint, ) @@ -78,7 +78,7 @@ def from_config( Configured WorkspaceConstraintChecker instance Example: - >>> from embodichain.lab.sim.workspace.configs import DimensionConstraint + >>> from embodichain.lab.sim.motion.workspace.configs import DimensionConstraint >>> config = DimensionConstraint( ... min_bounds=np.array([-1, -1, 0]), ... max_bounds=np.array([1, 1, 2]), @@ -222,8 +222,8 @@ def get_num_exclude_zones(self) -> int: 1. Basic usage - Setting workspace boundaries: ```python import numpy as np -from embodichain.lab.sim.workspace.configs import DimensionConstraint -from embodichain.lab.sim.workspace.constraints import WorkspaceConstraintChecker +from embodichain.lab.sim.motion.workspace.configs import DimensionConstraint +from embodichain.lab.sim.motion.workspace.constraints import WorkspaceConstraintChecker # Create via configuration (recommended approach) config = DimensionConstraint( @@ -270,7 +270,7 @@ def get_num_exclude_zones(self) -> int: 3. Application in robot workspace analysis: ```python # Usage in WorkspaceAnalyzer -from embodichain.lab.sim.workspace import WorkspaceAnalyzer +from embodichain.lab.sim.motion.workspace import WorkspaceAnalyzer # Configure work environment with complex obstacles config = DimensionConstraint( diff --git a/embodichain/lab/sim/workspace/metrics/__init__.py b/embodichain/lab/sim/motion/workspace/metrics/__init__.py similarity index 78% rename from embodichain/lab/sim/workspace/metrics/__init__.py rename to embodichain/lab/sim/motion/workspace/metrics/__init__.py index 3a6d31be2..93e1d7ca4 100644 --- a/embodichain/lab/sim/workspace/metrics/__init__.py +++ b/embodichain/lab/sim/motion/workspace/metrics/__init__.py @@ -18,16 +18,16 @@ Built-in metrics: reachability, manipulability, and density.""" -from embodichain.lab.sim.workspace.metrics.base_metric import ( +from embodichain.lab.sim.motion.workspace.metrics.base_metric import ( BaseMetric, ) -from embodichain.lab.sim.workspace.metrics.reachability_metric import ( +from embodichain.lab.sim.motion.workspace.metrics.reachability_metric import ( ReachabilityMetric, ) -from embodichain.lab.sim.workspace.metrics.manipulability_metric import ( +from embodichain.lab.sim.motion.workspace.metrics.manipulability_metric import ( ManipulabilityMetric, ) -from embodichain.lab.sim.workspace.metrics.density_metric import ( +from embodichain.lab.sim.motion.workspace.metrics.density_metric import ( DensityMetric, ) diff --git a/embodichain/lab/sim/workspace/metrics/base_metric.py b/embodichain/lab/sim/motion/workspace/metrics/base_metric.py similarity index 100% rename from embodichain/lab/sim/workspace/metrics/base_metric.py rename to embodichain/lab/sim/motion/workspace/metrics/base_metric.py diff --git a/embodichain/lab/sim/workspace/metrics/density_metric.py b/embodichain/lab/sim/motion/workspace/metrics/density_metric.py similarity index 97% rename from embodichain/lab/sim/workspace/metrics/density_metric.py rename to embodichain/lab/sim/motion/workspace/metrics/density_metric.py index a668ae3fc..ac5e96ddd 100644 --- a/embodichain/lab/sim/workspace/metrics/density_metric.py +++ b/embodichain/lab/sim/motion/workspace/metrics/density_metric.py @@ -16,10 +16,10 @@ from typing import Dict, Any import numpy as np -from embodichain.lab.sim.workspace.metrics.base_metric import ( +from embodichain.lab.sim.motion.workspace.metrics.base_metric import ( BaseMetric, ) -from embodichain.lab.sim.workspace.configs.metric_config import ( +from embodichain.lab.sim.motion.workspace.configs.metric_config import ( DensityConfig, ) diff --git a/embodichain/lab/sim/workspace/metrics/manipulability_metric.py b/embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py similarity index 97% rename from embodichain/lab/sim/workspace/metrics/manipulability_metric.py rename to embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py index 4d23989ed..b24b983a6 100644 --- a/embodichain/lab/sim/workspace/metrics/manipulability_metric.py +++ b/embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py @@ -16,10 +16,10 @@ from typing import Dict, Any import numpy as np -from embodichain.lab.sim.workspace.metrics.base_metric import ( +from embodichain.lab.sim.motion.workspace.metrics.base_metric import ( BaseMetric, ) -from embodichain.lab.sim.workspace.configs.metric_config import ( +from embodichain.lab.sim.motion.workspace.configs.metric_config import ( ManipulabilityConfig, ) diff --git a/embodichain/lab/sim/workspace/metrics/reachability_metric.py b/embodichain/lab/sim/motion/workspace/metrics/reachability_metric.py similarity index 96% rename from embodichain/lab/sim/workspace/metrics/reachability_metric.py rename to embodichain/lab/sim/motion/workspace/metrics/reachability_metric.py index 18d0df559..58cfee7c7 100644 --- a/embodichain/lab/sim/workspace/metrics/reachability_metric.py +++ b/embodichain/lab/sim/motion/workspace/metrics/reachability_metric.py @@ -16,10 +16,10 @@ from typing import Dict, Any import numpy as np -from embodichain.lab.sim.workspace.metrics.base_metric import ( +from embodichain.lab.sim.motion.workspace.metrics.base_metric import ( BaseMetric, ) -from embodichain.lab.sim.workspace.configs.metric_config import ( +from embodichain.lab.sim.motion.workspace.configs.metric_config import ( ReachabilityConfig, ) diff --git a/embodichain/lab/sim/workspace/runtime.py b/embodichain/lab/sim/motion/workspace/runtime.py similarity index 100% rename from embodichain/lab/sim/workspace/runtime.py rename to embodichain/lab/sim/motion/workspace/runtime.py diff --git a/embodichain/lab/sim/workspace/samplers/__init__.py b/embodichain/lab/sim/motion/workspace/samplers/__init__.py similarity index 75% rename from embodichain/lab/sim/workspace/samplers/__init__.py rename to embodichain/lab/sim/motion/workspace/samplers/__init__.py index 56061e720..cc4580adc 100644 --- a/embodichain/lab/sim/workspace/samplers/__init__.py +++ b/embodichain/lab/sim/motion/workspace/samplers/__init__.py @@ -21,28 +21,28 @@ Gaussian and importance strategies. """ -from embodichain.lab.sim.workspace.samplers.base_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.base_sampler import ( BaseSampler, ISampler, ) -from embodichain.lab.sim.workspace.samplers.iniform_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.iniform_sampler import ( UniformSampler, ) -from embodichain.lab.sim.workspace.samplers.random_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.random_sampler import ( RandomSampler, ) -from embodichain.lab.sim.workspace.samplers.halton_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.halton_sampler import ( HaltonSampler, ) -from embodichain.lab.sim.workspace.samplers.sobol_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.sobol_sampler import ( SobolSampler, ) -from embodichain.lab.sim.workspace.samplers.lhs_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.lhs_sampler import ( LatinHypercubeSampler, ) -from embodichain.lab.sim.workspace.samplers.sampler_factory import ( +from embodichain.lab.sim.motion.workspace.samplers.sampler_factory import ( SamplerFactory, create_sampler, ) @@ -51,7 +51,7 @@ # Note: PlaneSampler is intentionally not included here as it's a high-level # specialized sampler, not a base sampler. Import it directly when needed: -# from embodichain.lab.sim.workspace.samplers.plane_sampler import PlaneSampler +# from embodichain.lab.sim.motion.workspace.samplers.plane_sampler import PlaneSampler __all__ = [ # Base samplers and interfaces diff --git a/embodichain/lab/sim/workspace/samplers/base_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/base_sampler.py similarity index 100% rename from embodichain/lab/sim/workspace/samplers/base_sampler.py rename to embodichain/lab/sim/motion/workspace/samplers/base_sampler.py diff --git a/embodichain/lab/sim/workspace/samplers/gaussian_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/gaussian_sampler.py similarity index 98% rename from embodichain/lab/sim/workspace/samplers/gaussian_sampler.py rename to embodichain/lab/sim/motion/workspace/samplers/gaussian_sampler.py index 9e357d9fe..9c32ffddb 100644 --- a/embodichain/lab/sim/workspace/samplers/gaussian_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/gaussian_sampler.py @@ -18,11 +18,11 @@ import torch from typing import Union -from embodichain.lab.sim.workspace.configs.sampling_config import ( +from embodichain.lab.sim.motion.workspace.configs.sampling_config import ( SamplingStrategy, ) -from embodichain.lab.sim.workspace.samplers.base_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.base_sampler import ( BaseSampler, ) diff --git a/embodichain/lab/sim/workspace/samplers/halton_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/halton_sampler.py similarity index 98% rename from embodichain/lab/sim/workspace/samplers/halton_sampler.py rename to embodichain/lab/sim/motion/workspace/samplers/halton_sampler.py index 25cc5a12b..4a5d0dcc1 100644 --- a/embodichain/lab/sim/workspace/samplers/halton_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/halton_sampler.py @@ -18,10 +18,10 @@ import torch from typing import List, Union, TYPE_CHECKING -from embodichain.lab.sim.workspace.configs.sampling_config import ( +from embodichain.lab.sim.motion.workspace.configs.sampling_config import ( SamplingStrategy, ) -from embodichain.lab.sim.workspace.samplers.base_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.base_sampler import ( BaseSampler, ) diff --git a/embodichain/lab/sim/workspace/samplers/importance_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/importance_sampler.py similarity index 98% rename from embodichain/lab/sim/workspace/samplers/importance_sampler.py rename to embodichain/lab/sim/motion/workspace/samplers/importance_sampler.py index 425b24c1e..3eed34834 100644 --- a/embodichain/lab/sim/workspace/samplers/importance_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/importance_sampler.py @@ -18,10 +18,10 @@ import torch from typing import Callable, Union, TYPE_CHECKING -from embodichain.lab.sim.workspace.configs.sampling_config import ( +from embodichain.lab.sim.motion.workspace.configs.sampling_config import ( SamplingStrategy, ) -from embodichain.lab.sim.workspace.samplers.base_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.base_sampler import ( BaseSampler, ) diff --git a/embodichain/lab/sim/workspace/samplers/iniform_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/iniform_sampler.py similarity index 97% rename from embodichain/lab/sim/workspace/samplers/iniform_sampler.py rename to embodichain/lab/sim/motion/workspace/samplers/iniform_sampler.py index fb04a044b..46e40b793 100644 --- a/embodichain/lab/sim/workspace/samplers/iniform_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/iniform_sampler.py @@ -18,10 +18,10 @@ import torch from typing import Union, TYPE_CHECKING -from embodichain.lab.sim.workspace.configs.sampling_config import ( +from embodichain.lab.sim.motion.workspace.configs.sampling_config import ( SamplingStrategy, ) -from embodichain.lab.sim.workspace.samplers.base_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.base_sampler import ( BaseSampler, ) diff --git a/embodichain/lab/sim/workspace/samplers/lhs_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/lhs_sampler.py similarity index 98% rename from embodichain/lab/sim/workspace/samplers/lhs_sampler.py rename to embodichain/lab/sim/motion/workspace/samplers/lhs_sampler.py index 078351abc..f482506f1 100644 --- a/embodichain/lab/sim/workspace/samplers/lhs_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/lhs_sampler.py @@ -27,10 +27,10 @@ __all__ = ["LatinHypercubeSampler"] -from embodichain.lab.sim.workspace.configs.sampling_config import ( +from embodichain.lab.sim.motion.workspace.configs.sampling_config import ( SamplingStrategy, ) -from embodichain.lab.sim.workspace.samplers.base_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.base_sampler import ( BaseSampler, ) diff --git a/embodichain/lab/sim/workspace/samplers/random_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/random_sampler.py similarity index 96% rename from embodichain/lab/sim/workspace/samplers/random_sampler.py rename to embodichain/lab/sim/motion/workspace/samplers/random_sampler.py index a568826f2..e89b7a688 100644 --- a/embodichain/lab/sim/workspace/samplers/random_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/random_sampler.py @@ -18,10 +18,10 @@ import torch from typing import Union, TYPE_CHECKING -from embodichain.lab.sim.workspace.configs.sampling_config import ( +from embodichain.lab.sim.motion.workspace.configs.sampling_config import ( SamplingStrategy, ) -from embodichain.lab.sim.workspace.samplers.base_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.base_sampler import ( BaseSampler, ) diff --git a/embodichain/lab/sim/workspace/samplers/sampler_factory.py b/embodichain/lab/sim/motion/workspace/samplers/sampler_factory.py similarity index 92% rename from embodichain/lab/sim/workspace/samplers/sampler_factory.py rename to embodichain/lab/sim/motion/workspace/samplers/sampler_factory.py index ab19d179c..1964b0348 100644 --- a/embodichain/lab/sim/workspace/samplers/sampler_factory.py +++ b/embodichain/lab/sim/motion/workspace/samplers/sampler_factory.py @@ -21,31 +21,31 @@ import torch import numpy as np -from embodichain.lab.sim.workspace.configs.sampling_config import ( +from embodichain.lab.sim.motion.workspace.configs.sampling_config import ( SamplingStrategy, ) -from embodichain.lab.sim.workspace.samplers.base_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.base_sampler import ( BaseSampler, ) -from embodichain.lab.sim.workspace.samplers.iniform_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.iniform_sampler import ( UniformSampler, ) -from embodichain.lab.sim.workspace.samplers.random_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.random_sampler import ( RandomSampler, ) -from embodichain.lab.sim.workspace.samplers.halton_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.halton_sampler import ( HaltonSampler, ) -from embodichain.lab.sim.workspace.samplers.sobol_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.sobol_sampler import ( SobolSampler, ) -from embodichain.lab.sim.workspace.samplers.lhs_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.lhs_sampler import ( LatinHypercubeSampler, ) -from embodichain.lab.sim.workspace.samplers.importance_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.importance_sampler import ( ImportanceSampler, ) -from embodichain.lab.sim.workspace.samplers.gaussian_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.gaussian_sampler import ( GaussianSampler, ) diff --git a/embodichain/lab/sim/workspace/samplers/sobol_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/sobol_sampler.py similarity index 98% rename from embodichain/lab/sim/workspace/samplers/sobol_sampler.py rename to embodichain/lab/sim/motion/workspace/samplers/sobol_sampler.py index 6c8dcca42..f5ef64e3a 100644 --- a/embodichain/lab/sim/workspace/samplers/sobol_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/sobol_sampler.py @@ -25,10 +25,10 @@ except ImportError: SCIPY_AVAILABLE = False -from embodichain.lab.sim.workspace.configs.sampling_config import ( +from embodichain.lab.sim.motion.workspace.configs.sampling_config import ( SamplingStrategy, ) -from embodichain.lab.sim.workspace.samplers.base_sampler import ( +from embodichain.lab.sim.motion.workspace.samplers.base_sampler import ( BaseSampler, ) diff --git a/embodichain/lab/sim/workspace/visualizers/__init__.py b/embodichain/lab/sim/motion/workspace/visualizers/__init__.py similarity index 72% rename from embodichain/lab/sim/workspace/visualizers/__init__.py rename to embodichain/lab/sim/motion/workspace/visualizers/__init__.py index 836b46b45..c7692e14f 100644 --- a/embodichain/lab/sim/workspace/visualizers/__init__.py +++ b/embodichain/lab/sim/motion/workspace/visualizers/__init__.py @@ -19,31 +19,31 @@ Built-in visualizers: point-cloud, voxel, sphere, and axis, plus a ``VisualizerFactory`` and ``create_visualizer`` helper. """ -from embodichain.lab.sim.workspace.visualizers.base_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.base_visualizer import ( BaseVisualizer, IVisualizer, ) -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( VisualizationType, ) -from embodichain.lab.sim.workspace.visualizers.point_cloud_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.point_cloud_visualizer import ( PointCloudVisualizer, ) -from embodichain.lab.sim.workspace.visualizers.voxel_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.voxel_visualizer import ( VoxelVisualizer, ) -from embodichain.lab.sim.workspace.visualizers.sphere_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.sphere_visualizer import ( SphereVisualizer, ) -from embodichain.lab.sim.workspace.visualizers.axis_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.axis_visualizer import ( AxisVisualizer, ) -from embodichain.lab.sim.workspace.visualizers.visualizer_factory import ( +from embodichain.lab.sim.motion.workspace.visualizers.visualizer_factory import ( VisualizerFactory, create_visualizer, ) diff --git a/embodichain/lab/sim/workspace/visualizers/axis_visualizer.py b/embodichain/lab/sim/motion/workspace/visualizers/axis_visualizer.py similarity index 99% rename from embodichain/lab/sim/workspace/visualizers/axis_visualizer.py rename to embodichain/lab/sim/motion/workspace/visualizers/axis_visualizer.py index 35d274d9e..225ad84b2 100644 --- a/embodichain/lab/sim/workspace/visualizers/axis_visualizer.py +++ b/embodichain/lab/sim/motion/workspace/visualizers/axis_visualizer.py @@ -35,10 +35,10 @@ MATPLOTLIB_AVAILABLE = False from embodichain.utils import logger -from embodichain.lab.sim.workspace.visualizers.base_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.base_visualizer import ( BaseVisualizer, ) -from embodichain.lab.sim.workspace.configs.visualization_config import ( +from embodichain.lab.sim.motion.workspace.configs.visualization_config import ( VisualizationType, ) diff --git a/embodichain/lab/sim/workspace/visualizers/base_visualizer.py b/embodichain/lab/sim/motion/workspace/visualizers/base_visualizer.py similarity index 99% rename from embodichain/lab/sim/workspace/visualizers/base_visualizer.py rename to embodichain/lab/sim/motion/workspace/visualizers/base_visualizer.py index 8f7cb16f5..ce0884513 100644 --- a/embodichain/lab/sim/workspace/visualizers/base_visualizer.py +++ b/embodichain/lab/sim/motion/workspace/visualizers/base_visualizer.py @@ -36,7 +36,7 @@ MATPLOTLIB_AVAILABLE = False from embodichain.utils import logger -from embodichain.lab.sim.workspace.configs.visualization_config import ( +from embodichain.lab.sim.motion.workspace.configs.visualization_config import ( VisualizationConfig, ) diff --git a/embodichain/lab/sim/workspace/visualizers/point_cloud_visualizer.py b/embodichain/lab/sim/motion/workspace/visualizers/point_cloud_visualizer.py similarity index 98% rename from embodichain/lab/sim/workspace/visualizers/point_cloud_visualizer.py rename to embodichain/lab/sim/motion/workspace/visualizers/point_cloud_visualizer.py index f16469549..4c1b338e9 100644 --- a/embodichain/lab/sim/workspace/visualizers/point_cloud_visualizer.py +++ b/embodichain/lab/sim/motion/workspace/visualizers/point_cloud_visualizer.py @@ -21,10 +21,10 @@ from typing import Union, Any, Dict from pathlib import Path -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( VisualizationType, ) -from embodichain.lab.sim.workspace.visualizers.base_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.base_visualizer import ( BaseVisualizer, OPEN3D_AVAILABLE, ) diff --git a/embodichain/lab/sim/workspace/visualizers/sphere_visualizer.py b/embodichain/lab/sim/motion/workspace/visualizers/sphere_visualizer.py similarity index 98% rename from embodichain/lab/sim/workspace/visualizers/sphere_visualizer.py rename to embodichain/lab/sim/motion/workspace/visualizers/sphere_visualizer.py index ef6221826..c978b3472 100644 --- a/embodichain/lab/sim/workspace/visualizers/sphere_visualizer.py +++ b/embodichain/lab/sim/motion/workspace/visualizers/sphere_visualizer.py @@ -19,11 +19,11 @@ from typing import Union, Any, Dict from pathlib import Path -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( VisualizationType, ) -from embodichain.lab.sim.workspace.visualizers.base_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.base_visualizer import ( BaseVisualizer, OPEN3D_AVAILABLE, ) diff --git a/embodichain/lab/sim/workspace/visualizers/visualizer_factory.py b/embodichain/lab/sim/motion/workspace/visualizers/visualizer_factory.py similarity index 95% rename from embodichain/lab/sim/workspace/visualizers/visualizer_factory.py rename to embodichain/lab/sim/motion/workspace/visualizers/visualizer_factory.py index ecb177d58..1a95ceccc 100644 --- a/embodichain/lab/sim/workspace/visualizers/visualizer_factory.py +++ b/embodichain/lab/sim/motion/workspace/visualizers/visualizer_factory.py @@ -18,27 +18,27 @@ from typing import Dict, Type, Any from threading import Lock -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( VisualizationType, ) -from embodichain.lab.sim.workspace.visualizers.base_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.base_visualizer import ( BaseVisualizer, ) -from embodichain.lab.sim.workspace.visualizers.point_cloud_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.point_cloud_visualizer import ( PointCloudVisualizer, ) -from embodichain.lab.sim.workspace.visualizers.voxel_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.voxel_visualizer import ( VoxelVisualizer, ) -from embodichain.lab.sim.workspace.visualizers.sphere_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.sphere_visualizer import ( SphereVisualizer, ) -from embodichain.lab.sim.workspace.visualizers.axis_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.axis_visualizer import ( AxisVisualizer, ) diff --git a/embodichain/lab/sim/workspace/visualizers/voxel_visualizer.py b/embodichain/lab/sim/motion/workspace/visualizers/voxel_visualizer.py similarity index 98% rename from embodichain/lab/sim/workspace/visualizers/voxel_visualizer.py rename to embodichain/lab/sim/motion/workspace/visualizers/voxel_visualizer.py index 352aa521e..2271c66f4 100644 --- a/embodichain/lab/sim/workspace/visualizers/voxel_visualizer.py +++ b/embodichain/lab/sim/motion/workspace/visualizers/voxel_visualizer.py @@ -19,11 +19,11 @@ from typing import Union, Any, Dict from pathlib import Path -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( VisualizationType, ) -from embodichain.lab.sim.workspace.visualizers.base_visualizer import ( +from embodichain.lab.sim.motion.workspace.visualizers.base_visualizer import ( BaseVisualizer, OPEN3D_AVAILABLE, ) diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index 88f70be7c..a29380787 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -868,7 +868,7 @@ def root_link_name(self) -> str: Returns: str: The name of the root link. """ - return self.entities[0].get_root_link_name() + return self._entities[0].get_root_link_name() @cached_property def joint_names(self) -> List[str]: diff --git a/embodichain/lab/sim/objects/gizmo.py b/embodichain/lab/sim/objects/gizmo.py index 92b07d9e9..efd18d187 100644 --- a/embodichain/lab/sim/objects/gizmo.py +++ b/embodichain/lab/sim/objects/gizmo.py @@ -41,7 +41,7 @@ if TYPE_CHECKING: from dexsim.engine import GizmoController from dexsim.kit.ik import IKGizmoController, NewtonChainIK - from embodichain.lab.sim.solvers import BaseSolver + from embodichain.lab.sim.motion.solvers import BaseSolver __all__ = ["Gizmo", "GizmoCfg", "create_robot_ik_gizmo_controller"] diff --git a/embodichain/lab/sim/objects/robot.py b/embodichain/lab/sim/objects/robot.py index 9de9599e4..37aef4063 100644 --- a/embodichain/lab/sim/objects/robot.py +++ b/embodichain/lab/sim/objects/robot.py @@ -25,10 +25,10 @@ from dexsim.engine import Articulation as _Articulation from embodichain.lab.sim.cfg import RobotCfg, RobotWorkspaceCfg -from embodichain.lab.sim.solvers import SolverCfg, BaseSolver +from embodichain.lab.sim.motion.solvers import SolverCfg, BaseSolver from embodichain.lab.sim.objects import Articulation from embodichain.lab.sim.utility.tensor import to_tensor -from embodichain.lab.sim.workspace.runtime import ( +from embodichain.lab.sim.motion.workspace.runtime import ( RobotWorkspace, WorkspaceSample, ) diff --git a/embodichain/lab/sim/robots/cobotmagic.py b/embodichain/lab/sim/robots/cobotmagic.py index 66041772d..e6bd9b9f1 100644 --- a/embodichain/lab/sim/robots/cobotmagic.py +++ b/embodichain/lab/sim/robots/cobotmagic.py @@ -27,7 +27,7 @@ JointDrivePropertiesCfg, RigidBodyAttributesCfg, ) -from embodichain.lab.sim.solvers import SolverCfg, OPWSolverCfg +from embodichain.lab.sim.motion.solvers import SolverCfg, OPWSolverCfg from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg from embodichain.data import get_data_path from embodichain.utils import configclass diff --git a/embodichain/lab/sim/robots/dexforce_w1/cfg.py b/embodichain/lab/sim/robots/dexforce_w1/cfg.py index 2cea0df17..87eb31bb3 100644 --- a/embodichain/lab/sim/robots/dexforce_w1/cfg.py +++ b/embodichain/lab/sim/robots/dexforce_w1/cfg.py @@ -217,7 +217,7 @@ def _build_defaults(self, init_dict: dict | None = None) -> None: def _build_default_solver_cfg(self): """Build the version-matched default SRS solver configuration.""" - from embodichain.lab.sim.solvers import SRSSolverCfg + from embodichain.lab.sim.motion.solvers import SRSSolverCfg from embodichain.lab.sim.robots.dexforce_w1.params import ( W1ArmKineParams, ) diff --git a/embodichain/lab/sim/robots/dual_arm.py b/embodichain/lab/sim/robots/dual_arm.py index 6ddd63b28..aa14fcbf3 100644 --- a/embodichain/lab/sim/robots/dual_arm.py +++ b/embodichain/lab/sim/robots/dual_arm.py @@ -53,7 +53,7 @@ RobotCfg, URDFCfg, ) -from embodichain.lab.sim.solvers import SolverCfg +from embodichain.lab.sim.motion.solvers import SolverCfg from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg from embodichain.lab.sim.robots.franka_panda import FrankaPandaCfg from embodichain.lab.sim.robots.ur_robot import URRobotCfg diff --git a/embodichain/lab/sim/robots/franka_panda.py b/embodichain/lab/sim/robots/franka_panda.py index f80aeea96..e60b1cc3b 100644 --- a/embodichain/lab/sim/robots/franka_panda.py +++ b/embodichain/lab/sim/robots/franka_panda.py @@ -28,7 +28,7 @@ RobotCfg, URDFCfg, ) -from embodichain.lab.sim.solvers import PytorchSolverCfg +from embodichain.lab.sim.motion.solvers import PytorchSolverCfg from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg from embodichain.utils import configclass @@ -63,7 +63,7 @@ class FrankaPandaCfg(RobotCfg): The PandaWithHand URDF includes both the 7-DOF arm and the parallel-jaw gripper in a single file. The solver defaults to - :class:`~embodichain.lab.sim.solvers.PytorchSolverCfg`. + :class:`~embodichain.lab.sim.motion.solvers.PytorchSolverCfg`. Example: diff --git a/embodichain/lab/sim/robots/ur_robot.py b/embodichain/lab/sim/robots/ur_robot.py index b6a4f9135..6239d35d4 100644 --- a/embodichain/lab/sim/robots/ur_robot.py +++ b/embodichain/lab/sim/robots/ur_robot.py @@ -25,7 +25,7 @@ JointDrivePropertiesCfg, RigidBodyAttributesCfg, ) -from embodichain.lab.sim.solvers import URSolverCfg +from embodichain.lab.sim.motion.solvers import URSolverCfg from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg from embodichain.data import get_data_path from embodichain.utils import configclass @@ -64,7 +64,7 @@ class URRobotCfg(RobotCfg): One config class covers UR3 / UR3e / UR5 / UR5e / UR10 / UR10e, selected via ``robot_type``. The kinematic (DH) parameters are owned by - :class:`~embodichain.lab.sim.solvers.URSolverCfg`; this config owns the URDF, + :class:`~embodichain.lab.sim.motion.solvers.URSolverCfg`; this config owns the URDF, control parts, drive properties and rigid-body attributes. Example: diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 1bc8fdd9d..5f405753d 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -832,6 +832,16 @@ def render_camera_group(self, group_ids: list[int]) -> None: self._world.render_camera_group(group_ids) + @property + def simulation_time(self) -> float: + """Seconds successfully advanced through :meth:`update` since construction. + + Manual-step hosts can use differences of this clock for observation + timestamps. Direct backend updates performed outside this manager, + including some object preparation setters, are not included. + """ + return self._visualization_sim_time + def update(self, physics_dt: float | None = None, step: int = 10) -> None: """Advance physics explicitly and publish the resulting simulation state. @@ -3280,7 +3290,14 @@ def reset_objects_state( Args: env_ids (Sequence[int] | None): The environment IDs to reset. If None, reset all environments. excluded_uids (Sequence[str] | None): List of asset UIDs to exclude from resetting. If None, reset all assets. + + Raises: + RuntimeError: If a fixed-scene generation host owns this batch. """ + if getattr(self, "_trajectory_generation_owner", None) is not None: + raise RuntimeError( + "Normal scene reset is disabled while a generation host owns the batch." + ) excluded_uids = set(excluded_uids) if excluded_uids is not None else set() for uid, robot in self._robots.items(): if uid not in excluded_uids: diff --git a/embodichain/lab/sim/utility/action_utils.py b/embodichain/lab/sim/utility/action_utils.py index c528f8a12..351742c04 100644 --- a/embodichain/lab/sim/utility/action_utils.py +++ b/embodichain/lab/sim/utility/action_utils.py @@ -22,7 +22,7 @@ import torch import warp as wp -from embodichain.lab.sim.solvers.base_solver import BaseSolver +from embodichain.lab.sim.motion.solvers.base_solver import BaseSolver from embodichain.utils.device_utils import standardize_device_string from embodichain.utils.utility import inv_transform from embodichain.utils.warp import ( diff --git a/embodichain/lab/sim/utility/atom_action_utils.py b/embodichain/lab/sim/utility/atom_action_utils.py index 4525d5dd6..8898f0d73 100644 --- a/embodichain/lab/sim/utility/atom_action_utils.py +++ b/embodichain/lab/sim/utility/atom_action_utils.py @@ -21,7 +21,7 @@ from embodichain.utils.logger import log_error, log_warning from embodichain.lab.gym.utils.misc import mul_linear_expand -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( MoveType, PlanState, MotionGenerator, diff --git a/embodichain/lab/sim/utility/cfg_utils.py b/embodichain/lab/sim/utility/cfg_utils.py index 51ce7d028..722042513 100644 --- a/embodichain/lab/sim/utility/cfg_utils.py +++ b/embodichain/lab/sim/utility/cfg_utils.py @@ -15,7 +15,7 @@ # ---------------------------------------------------------------------------- from embodichain.lab.sim.cfg import RobotCfg -from embodichain.lab.sim.solvers import SolverCfg +from embodichain.lab.sim.motion.solvers import SolverCfg from embodichain.utils import logger diff --git a/embodichain/lab/task_program/integrations/simulation/environment.py b/embodichain/lab/task_program/integrations/simulation/environment.py index bb0d98343..bf0e20bfd 100644 --- a/embodichain/lab/task_program/integrations/simulation/environment.py +++ b/embodichain/lab/task_program/integrations/simulation/environment.py @@ -50,7 +50,7 @@ SceneSnapshot, TaskState, ) -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( BasePlannerCfg, MotionGenCfg, MotionGenerator, diff --git a/embodichain/lab/task_program/integrations/simulation/parallel_safety.py b/embodichain/lab/task_program/integrations/simulation/parallel_safety.py index 576b7a0c6..bbe1adfec 100644 --- a/embodichain/lab/task_program/integrations/simulation/parallel_safety.py +++ b/embodichain/lab/task_program/integrations/simulation/parallel_safety.py @@ -31,7 +31,7 @@ JointPositionTarget, RuntimeCommandFrame, ) -from embodichain.lab.sim.planners import CuroboPlanner, MotionGenerator +from embodichain.lab.sim.motion.planners import CuroboPlanner, MotionGenerator from embodichain.lab.task_program.semantics import ( RegistrySceneProvider, SceneRegistry, diff --git a/embodichain/lab/task_program/semantics/scene.py b/embodichain/lab/task_program/semantics/scene.py index 611395c07..732ef86e2 100644 --- a/embodichain/lab/task_program/semantics/scene.py +++ b/embodichain/lab/task_program/semantics/scene.py @@ -41,7 +41,7 @@ from .effects import EffectEvidenceAddress if TYPE_CHECKING: - from embodichain.lab.sim.planners import MotionGenerator + from embodichain.lab.sim.motion.planners import MotionGenerator from embodichain.lab.sim.sim_manager import SimulationManager diff --git a/embodichain/lab/trajectory_generation/__init__.py b/embodichain/lab/trajectory_generation/__init__.py new file mode 100644 index 000000000..3818cfd8e --- /dev/null +++ b/embodichain/lab/trajectory_generation/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__ = [] diff --git a/embodichain/lab/trajectory_generation/execution.py b/embodichain/lab/trajectory_generation/execution.py new file mode 100644 index 000000000..118b27785 --- /dev/null +++ b/embodichain/lab/trajectory_generation/execution.py @@ -0,0 +1,891 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Full-batch qpos rollouts with actual controller labels and causal evidence.""" + +from __future__ import annotations + +import math +import json +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field, fields, replace +from types import MappingProxyType +from typing import Any + +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CandidateIdentity, + CandidateTrajectoryBatch, + ExpertEpisode, + ValidationCheck, + ValidationResult, +) +from .initial_state import FixedSceneHost, PreparedBatch +from .integrations.contact import PickUpMotionValidator + +__all__ = ["QposRolloutExecutor"] + +_METADATA_RESERVE = 64 * 1024 +_EXECUTOR_METADATA_LIMIT = _METADATA_RESERVE // 2 + + +def _failure_text(error: Exception) -> str: + """Retain the concrete cause hidden by an execution-loop wrapper.""" + seen = {id(error)} + while error.__cause__ is not None and id(error.__cause__) not in seen: + error = error.__cause__ + seen.add(id(error)) + return f"{type(error).__name__}: {error}"[:512] + + +def _metadata_bytes(value: object) -> int: + def encode(item: object) -> dict: + if isinstance(item, Mapping): + return dict(item) + return {entry.name: getattr(item, entry.name) for entry in fields(item)} + + size = 0 + for chunk in json.JSONEncoder(default=encode, ensure_ascii=False).iterencode(value): + size += len(chunk.encode()) + if size > _EXECUTOR_METADATA_LIMIT: + break + return size + + +@dataclass +class _Row: + candidate: CandidateTrajectoryBatch + steps: int + observations: list[dict[str, torch.Tensor]] = field(default_factory=list) + actions: list[torch.Tensor] = field(default_factory=list) + timestamps: list[float] = field(default_factory=lambda: [0.0]) + started: bool = False + failure: str | None = None + fixed_world: bool = True + + +def _flatten_observation(value: Mapping, prefix: str = "") -> dict[str, torch.Tensor]: + """Preserve every nested tensor when no explicit Gym encoder is supplied.""" + result = {} + for key, item in value.items(): + if not isinstance(key, str) or not key or "." in key: + raise ValueError( + "Default Gym observation keys must be nonempty and contain no dots." + ) + name = f"{prefix}.{key}" if prefix else str(key) + if isinstance(item, torch.Tensor): + result[name] = item + elif isinstance(item, Mapping): + result.update(_flatten_observation(item, name)) + else: + raise ValueError( + f"Observation field {name!r} requires an explicit encoder." + ) + return result + + +class QposRolloutExecutor: + """Execute one candidate per physical row through Gym or a pure simulator. + + Each candidate contains one initial sample followed by controller targets; + ``valid_length - 1`` commands are executed on the fixed control clock. + Missing, finished, and terminated rows receive safe holds while remaining + rows advance in a single batched physics update. Holds and padding never + enter returned episodes. Gym execution uses the normal demonstration loop + and ``ControllerAction``; policy preprocessing is skipped once and normal + postprocessing remains enabled. Actual targets are captured before physics. + + The executor supports changes only to the host's active joints. Inactive + or mimic coordinates must remain constant unless a shared PickUp validator + models their complete geometry and physical motion. Task + validation runs at the final batch boundary; early trajectory interruption, + target disagreement, and observation/physics errors always reject evidence. + A submitted command without a complete observation transition still calls + ``on_started``, but cannot produce a fabricated training transition. + + Args: + host: Exclusive fixed-scene host; prepare it before calling execute. + control_dt: Control period, an integer multiple of profile physics_dt. + observe: Required pure-simulator observation callback, returning full + batch tensors. It is never used for Gym execution. + encode_observation: Optional Gym observation encoder. The default + flattens every tensor field. Uses prepared/step observations only. + validator: Trusted final task validator for one candidate and its owned + ``observations, actions, timestamps``. Must include ``validator_id``. + validator_id: Required task validation check identifier. + validation_profile_id: Registered execution validation profile identity. + max_episode_bytes: Maximum tensor plus metadata payload per row. Before + commands, reserve 64 KiB for bounded metadata in addition to tensors; + half remains available for downstream validation and storage schema. + contact_validator: Optional PickUp validator sharing this simulator and + robot. Adds physical-substep contact gates and measured object/TCP + observations; permits its target and observed mimic joints to move. + Supported only with pure simulation and the matching Runner planner. + """ + + def __init__( + self, + host: FixedSceneHost, + *, + control_dt: float, + observe: Callable[[], Mapping[str, torch.Tensor]] | None = None, + encode_observation: Callable[[Any], Mapping[str, torch.Tensor]] | None = None, + validator: Callable[ + [ + CandidateTrajectoryBatch, + Mapping[str, torch.Tensor], + torch.Tensor, + torch.Tensor, + ], + ValidationResult, + ], + validator_id: str = "task_success", + validation_profile_id: str = "verified_motion", + max_episode_bytes: int = 256 * 1024 * 1024, + contact_validator: PickUpMotionValidator | None = None, + ) -> None: + if ( + isinstance(control_dt, bool) + or not isinstance(control_dt, (int, float)) + or not math.isfinite(control_dt) + or control_dt <= 0 + ): + raise ValueError("control_dt must be finite and positive.") + if not callable(validator) or (host.env is None and not callable(observe)): + raise ValueError( + "Provide a task validator and a pure-sim observation callback when needed." + ) + if encode_observation is not None and not callable(encode_observation): + raise TypeError("encode_observation must be callable.") + if type(max_episode_bytes) is not int or max_episode_bytes <= 0: + raise ValueError("max_episode_bytes must be a positive integer.") + for name, value in ( + ("validator_id", validator_id), + ("validation_profile_id", validation_profile_id), + ): + if not isinstance(value, str) or not value or len(value.encode()) > 1024: + raise ValueError(f"{name} must be nonempty text of at most 1024 bytes.") + if validator_id in {"execution_complete", "fixed_collision_world"}: + raise ValueError("validator_id cannot replace an execution check.") + ratio = control_dt / host.profile.physics_dt + if round(ratio) < 1 or not math.isclose( + ratio, round(ratio), rel_tol=0, abs_tol=1e-8 + ): + raise ValueError( + "control_dt must contain an integer number of physics steps." + ) + if host.env is not None and not math.isclose( + control_dt, host.env.step_dt, rel_tol=0, abs_tol=1e-8 + ): + raise ValueError("control_dt must match the Gym control clock.") + self.host = host + self.control_dt = float(control_dt) + self.observe = observe + self.encode_observation = encode_observation or _flatten_observation + self.validator = validator + self.validator_id = validator_id + self.validation_profile_id = validation_profile_id + self.max_episode_bytes = max_episode_bytes + if contact_validator is not None and ( + not isinstance(contact_validator, PickUpMotionValidator) + or host.env is not None + or contact_validator.robot is not host.adapter.robot + or contact_validator.sim is not host.adapter.sim + ): + raise ValueError( + "Contact validation requires the same pure-sim host robot." + ) + self.contact_validator = contact_validator + self._substeps = round(ratio) + self._running = False + self._last_failures: Mapping[str, str] = MappingProxyType({}) + + @property + def last_failures(self) -> Mapping[str, str]: + """Read failure reasons from the last batch, including zero-frame rows. + + Cleared when a new serialized ``execute`` call begins. The owned, + immutable mapping contains at most one entry per candidate, each + limited to 512 characters. Failed physics/observation transitions retain + their concrete exception cause even when no episode can be returned. + """ + return self._last_failures + + def _observations( + self, raw: Any = None, schema: dict | None = None + ) -> dict[str, torch.Tensor]: + values = ( + self.observe() if self.host.env is None else self.encode_observation(raw) + ) + if not isinstance(values, Mapping) or not values: + raise ValueError( + "Observation providers must return a nonempty tensor mapping." + ) + values = dict(values) + if self.contact_validator is not None: + physical = self.contact_validator.observations() + if any( + key in values and not torch.equal(values[key], value) + for key, value in physical.items() + ): + raise ValueError( + "Contact observations must contain actual simulator poses." + ) + values.update(physical) + joints = self.host.adapter.robot.get_qpos() + if "joint_positions" in values and not torch.equal( + values["joint_positions"], joints + ): + raise ValueError( + "joint_positions must contain actual full-joint measurements." + ) + values["joint_positions"] = joints + batch = self.host.adapter.sim.num_envs + if joints.shape != (batch, len(self.host.adapter.robot.joint_names)): + raise ValueError( + "Measured joint positions do not match the full joint layout." + ) + for key, value in values.items(): + if ( + not isinstance(key, str) + or not key + or not isinstance(value, torch.Tensor) + or value.ndim < 1 + or value.shape[0] != batch + or value.is_complex() + ): + raise ValueError(f"Invalid full-batch observation field {key!r}.") + if ( + sum(value.numel() * value.element_size() for value in values.values()) + > batch * self.max_episode_bytes + ): + raise ValueError( + "Observation batch exceeds the payload budget before ownership copy." + ) + if ( + schema is not None + and {key: (value.shape, value.dtype) for key, value in values.items()} + != schema + ): + raise ValueError("Observation schema changed during the rollout.") + if not all(bool(torch.isfinite(value).all()) for value in values.values()): + raise ValueError("Observations must contain finite real values.") + return {key: value.detach().cpu().clone() for key, value in values.items()} + + def _active_joints(self) -> tuple[int, ...]: + source = self.host.env if self.host.env is not None else self.host.adapter.robot + indices = tuple(int(value) for value in source.active_joint_ids) + dof = len(self.host.adapter.robot.joint_names) + if ( + not indices + or len(set(indices)) != len(indices) + or any(index < 0 or index >= dof for index in indices) + or not set(indices).issubset(self.host.adapter.robot.active_joint_ids) + ): + raise ValueError("The host must declare unique active joint indices.") + return indices + + def _preflight( + self, + binding: PreparedBatch, + candidates: Sequence[CandidateTrajectoryBatch | None], + episode_ids: Mapping[str, tuple[str, str]], + initial: dict[str, torch.Tensor], + joints: tuple[int, ...], + ) -> list[_Row | None]: + batch = self.host.adapter.sim.num_envs + if len(candidates) != batch: + raise ValueError( + "Candidates must cover every physical row, using None for idle rows." + ) + if not self.host.verify_initial(binding).accepted: + raise ValueError( + "The physical batch no longer matches its prepared initial state." + ) + robot = self.host.adapter.robot + inactive = [ + index + for index in range(len(robot.joint_names)) + if index not in joints + and ( + self.contact_validator is None + or index not in self.contact_validator.mimic_ids + ) + ] + limits = robot.get_qpos_limits().detach().cpu() + rows = [] + seen = set() + for index, candidate in enumerate(candidates): + if candidate is None: + rows.append(None) + continue + if ( + not isinstance(candidate, CandidateTrajectoryBatch) + or len(candidate.identities) != 1 + ): + raise ValueError("Each physical row requires a single owned candidate.") + candidate = candidate.row(0) + identity = candidate.identities[0] + if identity.candidate_id in seen: + raise ValueError("A candidate cannot execute in two rows at once.") + seen.add(identity.candidate_id) + if ( + identity.candidate_id not in episode_ids + or len(episode_ids[identity.candidate_id]) != 2 + or any( + not isinstance(value, str) or not value + for value in episode_ids[identity.candidate_id] + ) + ): + raise ValueError( + "Every candidate requires stable episode and commit IDs." + ) + case = binding.cases[index] + if ( + identity.scene_case_id != case.scene_case_id + or identity.initial_state_id != case.initial_state_id + or candidate.joint_names != tuple(robot.joint_names) + ): + raise ValueError( + "Candidate case, initial state, or joint order does not match its slot." + ) + length = int(candidate.valid_length[0]) + if length < 2: + raise ValueError( + "A rollout requires an initial sample and at least one command." + ) + positions = candidate.positions[0, :length].detach().cpu() + intervals = candidate.dt[0, 1:length].detach().cpu().to(torch.float64) + if not torch.allclose( + intervals, + torch.full_like(intervals, self.control_dt), + atol=1e-8, + rtol=0, + ): + raise ValueError( + "Candidate timing must use the executor's control clock." + ) + if not torch.allclose( + positions[0], + initial["joint_positions"][index].to(positions.dtype), + atol=self.host.adapter.atol, + rtol=0, + ): + raise ValueError( + "Candidate first sample does not match the actual initial joints." + ) + if inactive and not torch.equal( + positions[:, inactive], positions[0, inactive].expand(length, -1) + ): + raise ValueError( + "Candidates cannot change inactive or mimic joint coordinates." + ) + if bool( + ( + (positions < limits[index, :, 0]) + | (positions > limits[index, :, 1]) + ).any() + ): + raise ValueError( + "Candidate joint targets exceed actual controller limits." + ) + bytes_per_obs = sum( + value[index].numel() * value.element_size() + for value in initial.values() + ) + size = ( + length * (bytes_per_obs + 8) + + (length - 1) + * initial["joint_positions"][index].numel() + * initial["joint_positions"].element_size() + + _METADATA_RESERVE + ) + if size > self.max_episode_bytes: + raise ValueError( + "Rollout would exceed max_episode_bytes before its first command." + ) + if ( + _metadata_bytes( + ( + identity, + episode_ids[identity.candidate_id], + candidate.phases, + tuple(initial), + ) + ) + > _EXECUTOR_METADATA_LIMIT + ): + raise ValueError( + "Candidate lineage or observation keys exceed the metadata budget." + ) + rows.append( + _Row( + candidate, + length - 1, + [{key: value[index].clone() for key, value in initial.items()}], + ) + ) + return rows + + def _hold(self, joints: tuple[int, ...]) -> None: + robot = self.host.adapter.robot + measured = robot.get_qpos()[:, joints] + if not bool(torch.isfinite(measured).all()): + raise RuntimeError("Cannot safely hold nonfinite measured joint positions.") + robot.set_qpos(measured, joint_ids=joints, target=True) + robot.set_qvel(torch.zeros_like(measured), joint_ids=joints, target=True) + robot.set_qf(torch.zeros_like(measured), joint_ids=joints) + + def _clear_hold_dynamics( + self, joints: tuple[int, ...], active: tuple[bool, ...] + ) -> None: + if all(active): + return + robot = self.host.adapter.robot + velocity = robot.get_qvel(target=True)[:, joints].clone() + effort = robot.get_qf()[:, joints].clone() + inactive = torch.tensor([not value for value in active], device=velocity.device) + velocity[inactive] = 0 + effort[inactive] = 0 + robot.set_qvel(velocity, joint_ids=joints, target=True) + robot.set_qf(effort, joint_ids=joints) + + def execute( + self, + binding: PreparedBatch, + candidates: Sequence[CandidateTrajectoryBatch | None], + episode_ids: Mapping[str, tuple[str, str]], + *, + on_started: Callable[[CandidateIdentity], None], + should_stop: Callable[[], bool] | None = None, + ) -> tuple[ExpertEpisode | None, ...]: + """Execute the full batch and freeze each row at its final observed step. + + Invalid inputs raise before commands. Runtime failures yield rejected + episodes for rows with complete transitions and ``None`` otherwise. + ``on_started`` fires after the first submitted command, even when its + subsequent physics or observation fails. Inability to install the final + safe hold raises and must stop the caller's collection job. + """ + if self._running: + raise RuntimeError("QposRolloutExecutor requires serialized execution.") + self._last_failures = MappingProxyType({}) + if not callable(on_started) or ( + should_stop is not None and not callable(should_stop) + ): + raise TypeError("Execution callbacks must be callable.") + self.host.assert_current(binding) + initial = self._observations(self.host.initial_observation(binding)) + joints = self._active_joints() + rows = self._preflight(binding, candidates, episode_ids, initial, joints) + if not any(row is not None for row in rows): + return tuple(None for _ in rows) + schema = {key: (value.shape, value.dtype) for key, value in initial.items()} + robot, sim, env = self.host.adapter.robot, self.host.adapter.sim, self.host.env + snapshots = self.host.snapshots(binding) + if self.contact_validator is not None: + self.contact_validator.begin_rollout(candidates, snapshots) + initial_time = float(sim.simulation_time) + if not math.isfinite(initial_time): + raise ValueError("The simulator clock must be finite.") + last_time = initial_time + pending: torch.Tensor | None = None + current_mask = tuple(False for _ in rows) + expected_command: torch.Tensor | None = None + self._running = True + stop_requested = False + + def stop() -> bool: + nonlocal stop_requested + requested = should_stop is not None and bool(should_stop()) + stop_requested |= requested + return requested + + def command_submitted() -> None: + nonlocal pending, current_mask + if env is not None: + current_mask = tuple(bool(value) for value in env._demo_active_mask) + callback_error = None + for index, row in enumerate(rows): + if row is not None and current_mask[index] and not row.started: + row.started = True + try: + on_started(row.candidate.identities[0]) + except Exception as error: + callback_error = error + pending = robot.get_qpos(target=True).detach().cpu().clone() + if ( + pending.dtype != initial["joint_positions"].dtype + or pending.shape != initial["joint_positions"].shape + or not bool(torch.isfinite(pending).all()) + ): + raise ValueError( + "Actual controller targets have invalid shape or values." + ) + for index, row in enumerate(rows): + if ( + row is not None + and current_mask[index] + and not torch.allclose( + pending[index, joints], + expected_command[index, joints].cpu().to(pending.dtype), + atol=self.host.adapter.atol, + rtol=0, + ) + ): + row.failure = "actual controller targets differ from the candidate" + if callback_error is not None: + raise callback_error + + def transition(raw: Any = None, mask: tuple[bool, ...] | None = None) -> None: + nonlocal pending, last_time + self.host.assert_current(binding) + observed = self._observations(raw, schema) + now = float(sim.simulation_time) + if not math.isfinite(now) or not math.isclose( + now - last_time, self.control_dt, rel_tol=0, abs_tol=1e-8 + ): + raise ValueError( + "Actual simulator time does not match the control clock." + ) + last_time = now + if pending is None: + raise RuntimeError( + "A transition has no observed controller submission." + ) + active = current_mask if mask is None else mask + root_pose = robot.get_local_pose(to_matrix=True).detach().cpu() + entity_poses = { + uid: sim.get_rigid_object(uid) + .get_local_pose(to_matrix=True) + .detach() + .cpu() + for uid in sim.get_rigid_object_uid_list() + } + for index, row in enumerate(rows): + if row is not None: + snapshot = snapshots[index] + row.fixed_world &= ( + set(entity_poses) == set(snapshot.entity_poses) + and torch.allclose( + root_pose[index], + snapshot.root_pose.cpu(), + atol=self.host.adapter.atol, + rtol=0, + ) + and all( + torch.allclose( + pose[index], + snapshot.entity_poses[uid].cpu(), + atol=self.host.adapter.atol, + rtol=0, + ) + for uid, pose in entity_poses.items() + if self.contact_validator is None + or uid not in self.contact_validator.dynamic_entity_ids + ) + ) + if row is not None and active[index] and len(row.actions) < row.steps: + inactive = [ + joint + for joint in range(len(robot.joint_names)) + if joint not in joints + and ( + self.contact_validator is None + or joint not in self.contact_validator.mimic_ids + ) + ] + if inactive and ( + not torch.allclose( + observed["joint_positions"][index, inactive], + initial["joint_positions"][index, inactive], + atol=self.host.adapter.atol, + rtol=0, + ) + or not torch.allclose( + pending[index, inactive], + initial["joint_positions"][index, inactive].to(pending), + atol=self.host.adapter.atol, + rtol=0, + ) + ): + row.failure = ( + "inactive or mimic joints changed during execution" + ) + row.actions.append(pending[index].clone()) + row.timestamps.append(now - initial_time) + row.observations.append( + {key: value[index].clone() for key, value in observed.items()} + ) + pending = None + + def actions(): + nonlocal expected_command, current_mask + for index in range(max(row.steps for row in rows if row is not None)): + self.host.assert_current(binding) + expected_command = robot.get_qpos().detach().clone() + current_mask = tuple( + row is not None and index < row.steps for row in rows + ) + hold_mask = ( + current_mask + if env is None + else tuple( + active and bool(env._demo_active_mask[slot]) + for slot, active in enumerate(current_mask) + ) + ) + self._clear_hold_dynamics(joints, hold_mask) + for slot, row in enumerate(rows): + if row is not None and current_mask[slot]: + expected_command[slot] = row.candidate.positions[ + 0, index + 1 + ].to(expected_command) + yield expected_command.clone() + + try: + if env is not None: + from embodichain.lab.gym.envs.demo import ( + DemoSegment, + execute_demo_episode, + ) + from embodichain.lab.gym.envs.types import ControllerAction + + def observe_step(result, mask) -> None: + transition(result[0], mask) + terminated, truncated, info = result[2], result[3], result[4] + for index, row in enumerate(rows): + if row is None or not mask[index]: + continue + if bool(truncated[index]) or bool( + info.get("fail", torch.zeros(len(rows), dtype=torch.bool))[ + index + ] + ): + row.failure = "environment failed or truncated" + elif bool(terminated[index]) and len(row.actions) < row.steps: + row.failure = ( + "environment terminated before the trajectory ended" + ) + + with env.observe_generation_commands(self.host, command_submitted): + execute_demo_episode( + env, + segments=DemoSegment( + actions=(ControllerAction(value) for value in actions()), + name="qpos_candidate", + failure_policy="row_independent", + ), + row_step_limits=tuple( + row.steps if row is not None else 0 for row in rows + ), + step_observer=observe_step, + should_stop=stop, + ) + else: + for command in actions(): + if stop(): + break + robot.set_qpos(command[:, joints], joint_ids=joints, target=True) + command_submitted() + if self.contact_validator is None: + sim.update(self.host.profile.physics_dt, self._substeps) + else: + sample_index = ( + max(len(row.actions) for row in rows if row is not None) + 1 + ) + for _ in range(self._substeps): + sim.update(self.host.profile.physics_dt, 1) + self.contact_validator.observe_substep( + sample_index, + current_mask, + physics_dt=self.host.profile.physics_dt, + ) + transition() + except Exception as error: + for row in rows: + if row is not None: + row.failure = _failure_text(error) + finally: + try: + self._hold(joints) + except Exception as error: + self._last_failures = MappingProxyType( + { + row.candidate.identities[ + 0 + ].candidate_id: f"Safe hold failed: {_failure_text(error)}"[ + :512 + ] + for row in rows + if row is not None + } + ) + self._running = False + raise + + for row in rows: + if row is not None and row.failure is None and len(row.actions) < row.steps: + row.failure = ( + "execution cancelled by should_stop" + if stop_requested + else "trajectory interrupted before completion" + ) + self._last_failures = MappingProxyType( + { + row.candidate.identities[0].candidate_id: row.failure[:512] + for row in rows + if row is not None and row.failure is not None + } + ) + try: + episodes = self._freeze(rows, episode_ids) + failures = dict(self._last_failures) + for episode in episodes: + if episode is not None and not episode.validation.accepted: + failure = next( + check + for check in episode.validation.checks + if check.status != "passed" + ) + failures[episode.identity.candidate_id] = ( + failure.detail or failure.check_id + )[:512] + self._last_failures = MappingProxyType(failures) + return episodes + finally: + self._running = False + + def _freeze( + self, + rows: list[_Row | None], + episode_ids: Mapping[str, tuple[str, str]], + ) -> tuple[ExpertEpisode | None, ...]: + episodes = [] + for row_index, row in enumerate(rows): + if row is None or not row.actions: + episodes.append(None) + continue + observations = { + key: torch.stack([sample[key] for sample in row.observations]) + for key in row.observations[0] + } + commands = torch.stack(row.actions) + timestamps = torch.tensor(row.timestamps, dtype=torch.float64) + complete = len(commands) == row.steps and row.failure is None + checks = [ + ValidationCheck( + "execution_complete", + "passed" if complete else "failed", + row.failure + or ("" if complete else "trajectory interrupted before completion"), + ), + ValidationCheck( + "fixed_collision_world", + "passed" if row.fixed_world else "failed", + ( + "" + if row.fixed_world + else "Robot root or obstacle poses changed during execution." + ), + ), + ] + if self.contact_validator is not None: + checks.extend( + self.contact_validator.rollout_validation(row_index).checks + ) + try: + validation = self.validator( + row.candidate, + {key: value.clone() for key, value in observations.items()}, + commands.clone(), + timestamps.clone(), + ) + if not isinstance(validation, ValidationResult): + raise TypeError("Task validator must return ValidationResult.") + if not any( + check.check_id == self.validator_id for check in validation.checks + ): + raise ValueError( + f"Task validator omitted required check {self.validator_id!r}." + ) + if any( + check.check_id in {"execution_complete", "fixed_collision_world"} + for check in validation.checks + ): + raise ValueError("Task validator cannot replace execution checks.") + checks.extend(validation.checks) + except Exception as error: + checks.append( + ValidationCheck(self.validator_id, "failed", str(error)[:512]) + ) + episode_id, commit_id = episode_ids[ + row.candidate.identities[0].candidate_id + ] + phases = tuple( + replace(phase, stop_index=min(phase.stop_index, len(commands) + 1)) + for phase in row.candidate.phases[0] + if phase.start_index < len(commands) + 1 + ) + metadata = { + "validation_profile_id": self.validation_profile_id, + "timing_source": "simulation_time", + "phase_annotations": "planned_sample_intervals", + "joint_names": row.candidate.joint_names, + "commanded_joint_indices": self._active_joints(), + } + if self.contact_validator is not None: + metadata["contact_profile"] = self.contact_validator.profile.to_dict() + metadata["collision_geometry"] = ( + "URDF collision convex hulls; cuboid world; implicit ground" + ) + if ( + _metadata_bytes( + ( + row.candidate.identities[0], + episode_id, + commit_id, + "qpos", + tuple(observations), + ValidationResult(tuple(checks)), + phases, + metadata, + ) + ) + > _EXECUTOR_METADATA_LIMIT + ): + checks = [ + ValidationCheck( + "execution_complete", + "failed", + "Execution validation metadata exceeds its reserved budget.", + ) + ] + episode = ExpertEpisode( + identity=row.candidate.identities[0], + observations=observations, + actions=commands, + timestamps=timestamps, + action_representation="qpos", + validation=ValidationResult(tuple(checks)), + episode_id=episode_id, + commit_id=commit_id, + phases=phases, + metadata=metadata, + ) + episodes.append(episode) + return tuple(episodes) diff --git a/embodichain/lab/trajectory_generation/initial_state.py b/embodichain/lab/trajectory_generation/initial_state.py new file mode 100644 index 000000000..26caa796b --- /dev/null +++ b/embodichain/lab/trajectory_generation/initial_state.py @@ -0,0 +1,398 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Exclusive full-batch preparation for fixed-scene trajectory collection.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import MISSING, dataclass +import math +from typing import TYPE_CHECKING +from uuid import uuid4 + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + MotionSnapshot, + SceneCase, + ValidationCheck, + ValidationResult, +) +from embodichain.utils import configclass + +if TYPE_CHECKING: + from embodichain.lab.gym.envs.embodied_env import EmbodiedEnv + from embodichain.lab.sim.types import EnvObs + from .integrations.sim import SimInitialState, SimInitialStateAdapter + +__all__ = ["InitialStateProfile", "PreparedBatch", "FixedSceneHost"] + + +@configclass +class InitialStateProfile: + """Trusted preparation callbacks for one task and controller configuration. + + ``prepare`` must reset all task/controller state not owned by the standard + Gym managers without randomizing the scene. ``signature`` must identify the + fixed physical, visual, sensor, and control conditions, excluding evolving + poses and joint positions. ``verify`` checks the provided task initial state + against every declared case after settling. These callables are supplied + by trusted Python integration code, never imported from job YAML strings. + + Args: + profile_id: Stable preparation profile identifier. + prepare: Deterministic task/controller preparation callback. + signature: Callback returning a nonempty fixed-condition signature. + verify: Task-specific initial-state checks for the complete batch. + physics_dt: Physics step duration in seconds; must match Gym when used. + settling_steps: Number of preparation physics steps, outside rollouts. + allowed_interval_events: Gym interval event names explicitly certified + to preserve the fixed conditions by this profile. + """ + + profile_id: str = MISSING + prepare: Callable[[], None] = MISSING + signature: Callable[[], str] = MISSING + verify: Callable[[tuple[SceneCase, ...]], ValidationResult] = MISSING + physics_dt: float = 0.01 + settling_steps: int = 0 + allowed_interval_events: tuple[str, ...] = () + + def __post_init__(self) -> None: + if not isinstance(self.profile_id, str) or not self.profile_id.strip(): + raise ValueError("profile_id must be a nonempty string") + for name in ("prepare", "signature", "verify"): + if not callable(getattr(self, name)): + raise ValueError(f"{name} must be a trusted callback") + if ( + isinstance(self.physics_dt, bool) + or not isinstance(self.physics_dt, (int, float)) + or not math.isfinite(self.physics_dt) + or self.physics_dt <= 0 + ): + raise ValueError("physics_dt must be positive and finite") + if type(self.settling_steps) is not int or self.settling_steps < 0: + raise ValueError("settling_steps must be a nonnegative integer") + if isinstance(self.allowed_interval_events, str): + raise ValueError("allowed_interval_events must be a sequence of names") + self.allowed_interval_events = tuple(self.allowed_interval_events) + if any( + not isinstance(name, str) or not name.strip() + for name in self.allowed_interval_events + ) or len(set(self.allowed_interval_events)) != len( + self.allowed_interval_events + ): + raise ValueError("allowed_interval_events must contain unique names") + + +@dataclass(frozen=True) +class PreparedBatch: + """Identity of a successfully prepared full batch. + + Cases are ordered by physical slot. A token is valid only for its owning + host and epoch; it does not make an old ActionPlan or runtime reusable. + """ + + host_id: str + epoch: int + cases: tuple[SceneCase, ...] + + +class FixedSceneHost: + """Own preparation of one complete simulator batch, optionally through Gym. + + Capture consumes a caller-provided scene. Restoring discards pending Gym + recordings, so the caller must freeze old episode evidence first. Both + paths run deterministic preparation, settling, and verification before + publishing a new binding. The owner must remain the sole simulation + stepper for the lifetime of this object. + + Args: + adapter: Physical state adapter for the complete simulator batch. + profile: Trusted task, controller, event, and fixed-condition policy. + env: Optional unwrapped Gym host using the same simulator and robot. + + Raises: + ValueError: If the host or preparation policy is incompatible. + RuntimeError: If another generation host already owns the batch. + """ + + def __init__( + self, + adapter: SimInitialStateAdapter, + profile: InitialStateProfile, + *, + env: EmbodiedEnv | None = None, + ) -> None: + self.adapter = adapter + self.profile = profile.copy() + self.env = env + self._host_id = uuid4().hex + self._epoch = 0 + self._binding: PreparedBatch | None = None + self._initial_state: SimInitialState | None = None + self._cases: tuple[SceneCase, ...] = () + self._condition_signature: str | None = None + self._prepared_observation: EnvObs | None = None + self._closed = False + if env is not None: + if env.sim is not adapter.sim or env.robot is not adapter.robot: + raise ValueError("Gym and physical adapter must own the same sim/robot") + if env.num_envs != adapter.sim.num_envs: + raise ValueError("generation requires the complete simulator batch") + if not math.isclose(env.physics_dt, profile.physics_dt, rel_tol=1e-9): + raise ValueError("profile physics_dt must match the Gym physics clock") + self._check_interval_events() + if getattr(adapter.sim, "_trajectory_generation_owner", None) is not None: + raise RuntimeError("simulator batch already has a generation owner") + self._structure_signature = adapter.signature() + if env is not None: + env.acquire_generation_lease(self) + adapter.sim._trajectory_generation_owner = self + + def acquire_case(self, cases: Sequence[SceneCase]) -> PreparedBatch: + """Prepare and capture all caller-provided initial states. + + Args: + cases: One immutable scene case per simulator row. + + Returns: + A binding published only after preparation and validation succeed. + + Raises: + ValueError: If cases do not cover the full batch. + RuntimeError: If this host already acquired its fixed cases. + """ + self._assert_owner() + if self._initial_state is not None: + raise RuntimeError("cases already acquired; use restore_initial") + values = tuple(cases) + if len(values) != self.adapter.sim.num_envs or any( + not isinstance(case, SceneCase) for case in values + ): + raise ValueError("provide one SceneCase for every simulator row") + self._cases = values + self._condition_signature = self._signature() + return self._prepare(capture=True) + + def restore_initial(self) -> PreparedBatch: + """Restore the owned initial batch and invalidate every earlier binding. + + Returns: + The new successfully prepared binding. + + Raises: + RuntimeError: If no initial batch exists or fixed conditions drifted. + """ + self._assert_owner() + if self._initial_state is None: + raise RuntimeError("acquire_case must succeed before restore_initial") + return self._prepare(capture=False) + + def assert_current(self, batch: PreparedBatch) -> None: + """Reject a foreign or obsolete execution binding before issuing commands. + + Args: + batch: The binding attached to the candidate execution. + + Raises: + RuntimeError: If ownership, epoch, or fixed conditions changed. + """ + self._assert_owner() + if batch is not self._binding or batch.epoch != self._current_epoch(): + raise RuntimeError("execution binding belongs to an obsolete host epoch") + self._check_conditions() + + def verify_initial(self, batch: PreparedBatch) -> ValidationResult: + """Recheck physical and task initial state for the current binding. + + Args: + batch: The current prepared binding, before rollout begins. + + Returns: + Physical restoration and task-specific verification checks. + """ + self.assert_current(batch) + return self._verify() + + def snapshots(self, batch: PreparedBatch) -> tuple[MotionSnapshot, ...]: + """Export owned planning inputs from the captured initial batch. + + Args: + batch: Current preparation binding for the declared cases. + + Returns: + One full-joint snapshot per physical row, in local arena coordinates. + """ + self.assert_current(batch) + state = self._initial_state + return tuple( + MotionSnapshot( + scene_case=case, + joint_names=state.joint_names, + joint_positions=state.robot["qpos"][row], + joint_velocities=state.robot["qvel"][row], + root_pose=state.robot["root_pose"][row], + entity_poses={ + uid: fields["pose"][row] + for uid, fields in state.rigid_objects.items() + }, + ) + for row, case in enumerate(batch.cases) + ) + + def initial_observation(self, batch: PreparedBatch) -> EnvObs | None: + """Copy the Gym observation published by preparation without querying again. + + Args: + batch: Current preparation binding. + + Returns: + The owned Gym first frame, or ``None`` for a pure simulator host. + """ + self.assert_current(batch) + return ( + None + if self._prepared_observation is None + else self._prepared_observation.clone() + ) + + def close(self) -> None: + """Release exclusive ownership and invalidate all execution bindings.""" + if self._closed: + return + self._assert_owner() + if self.env is not None: + self.env.release_generation_lease(self) + self.adapter.sim._trajectory_generation_owner = None + self._binding = None + self._prepared_observation = None + self._epoch += 1 + self._closed = True + + def __enter__(self) -> FixedSceneHost: + self._assert_owner() + return self + + def __exit__(self, *args: object) -> None: + self.close() + + def _assert_owner(self) -> None: + if ( + self._closed + or getattr(self.adapter.sim, "_trajectory_generation_owner", None) + is not self + ): + raise RuntimeError("generation host does not own the simulator batch") + + def _signature(self) -> str: + value = self.profile.signature() + if not isinstance(value, str) or not value.strip(): + raise ValueError("fixed-condition signature must be a nonempty string") + return value + + def _check_interval_events(self) -> None: + manager = getattr(self.env, "event_manager", None) + active = manager.active_functors if manager is not None else {} + unknown = set(active.get("interval", ())) - set( + self.profile.allowed_interval_events + ) + if unknown: + raise ValueError(f"uncertified interval events: {sorted(unknown)}") + + def _check_conditions(self) -> None: + self._check_interval_events() + if self.adapter.signature() != self._structure_signature: + raise RuntimeError("simulator topology or controller identity changed") + if self._signature() != self._condition_signature: + raise RuntimeError("fixed scene conditions changed after case acquisition") + + def _current_epoch(self) -> int: + return self.env.generation_epoch if self.env is not None else self._epoch + + def _settle(self) -> None: + if self.profile.settling_steps: + self.adapter.sim.update( + self.profile.physics_dt, self.profile.settling_steps + ) + + def _verify(self) -> ValidationResult: + self._check_conditions() + task_result = self.profile.verify(self._cases) + if not isinstance(task_result, ValidationResult): + raise TypeError("profile verify must return ValidationResult") + if not task_result.checks: + task_result = ValidationResult( + ( + ValidationCheck( + "initial_conditions", "failed", "profile rejected initial state" + ), + ) + ) + physical = ( + self.adapter.verify(self._initial_state) + if self._initial_state is not None + else ValidationResult(()) + ) + return ValidationResult((*physical.checks, *task_result.checks)) + + def _prepare(self, *, capture: bool) -> PreparedBatch: + self._binding = None + self._prepared_observation = None + self._epoch += 1 + + def prepare() -> None: + self._check_conditions() + self.profile.prepare() + + def restore() -> None: + if not capture: + self.adapter.restore(self._initial_state) + + def verify() -> ValidationResult: + result = self._verify() + if result.accepted and capture: + state = self.adapter.capture() + physical = self.adapter.verify(state) + result = ValidationResult((*physical.checks, *result.checks)) + if result.accepted: + self._initial_state = state + return result + + try: + if self.env is not None: + observation, _ = self.env.prepare_generation_episode( + self, + prepare=prepare, + restore=restore, + settle=self._settle, + verify=verify, + ) + self._prepared_observation = observation.clone() + else: + prepare() + restore() + self._settle() + result = verify() + if not result.accepted: + raise RuntimeError( + f"initial state verification failed: {result.checks}" + ) + except Exception: + if capture: + self._initial_state = None + raise + self._binding = PreparedBatch(self._host_id, self._current_epoch(), self._cases) + return self._binding diff --git a/embodichain/lab/trajectory_generation/integrations/__init__.py b/embodichain/lab/trajectory_generation/integrations/__init__.py new file mode 100644 index 000000000..3818cfd8e --- /dev/null +++ b/embodichain/lab/trajectory_generation/integrations/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__ = [] diff --git a/embodichain/lab/trajectory_generation/integrations/_collision.py b/embodichain/lab/trajectory_generation/integrations/_collision.py new file mode 100644 index 000000000..9e894902b --- /dev/null +++ b/embodichain/lab/trajectory_generation/integrations/_collision.py @@ -0,0 +1,189 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Conservative CPU collision geometry for the bounded PickUp integration.""" + +from __future__ import annotations + +from collections import defaultdict +from itertools import combinations +import numpy as np +import torch + +__all__ = [] + + +class _FullStateCollisionWorld: + """URDF collision convex hulls and explicitly supported physical cuboids. + + Never substitute render meshes for collision geometry. One convex hull per + link conservatively covers all its URDF collision shapes, including their + local transforms. Unsupported scene shapes/scales fail at construction. + """ + + def __init__(self, sim, robot): + import fcl + import trimesh + from yourdfpy import URDF + from embodichain.lab.sim.shapes import CubeCfg + + if tuple(robot.cfg.body_scale) != (1.0, 1.0, 1.0): + raise ValueError("PickUp collision validation requires an unscaled URDF") + self.fcl, self.robot = fcl, robot + urdf = URDF.load( + robot.cfg.fpath, + build_scene_graph=False, + load_meshes=False, + build_collision_scene_graph=True, + load_collision_meshes=True, + ) + scene = urdf.collision_scene + meshes = defaultdict(list) + for node in scene.graph.nodes_geometry: + parent = scene.graph.transforms.parents[node] + if parent not in urdf.link_map: + raise ValueError( + "Collision geometry is not directly owned by a URDF link" + ) + transform, name = scene.graph.get(frame_to=node, frame_from=parent) + mesh = scene.geometry[name].copy() + mesh.apply_transform(transform) + meshes[parent].append(mesh) + expected = {name for name, link in urdf.link_map.items() if link.collisions} + if not expected or set(meshes) != expected: + raise ValueError("Every URDF collision link must have loaded geometry") + self.links = tuple(meshes) + neighbors = defaultdict(set) + fixed_links = {urdf.base_link} + while True: + previous = len(fixed_links) + fixed_links.update( + j.child + for j in urdf.joint_map.values() + if j.type == "fixed" and j.parent in fixed_links + ) + if previous == len(fixed_links): + break + self.fixed_links = frozenset(fixed_links) + for joint in urdf.joint_map.values(): + neighbors[joint.parent].add(joint.child) + neighbors[joint.child].add(joint.parent) + # Match the existing cuRobo profile's two-hop structural exclusions. + self.adjacent_pairs = set() + for link in self.links: + nearby = neighbors[link] | set().union( + *(neighbors[n] for n in neighbors[link]) + ) + self.adjacent_pairs.update( + frozenset((link, other)) for other in nearby if other != link + ) + self.objects = {} + self.bounds = {} + for link, parts in meshes.items(): + hull = trimesh.util.concatenate(parts).convex_hull + faces = np.c_[np.full(len(hull.faces), 3), hull.faces].astype(np.int32) + geometry = fcl.Convex( + hull.vertices.astype(np.float64), len(faces), faces.ravel() + ) + self.objects[link] = fcl.CollisionObject(geometry) + self.bounds[link] = hull.bounds + self.entity_ids = tuple(sim.get_rigid_object_uid_list()) + if set(self.links) & set(self.entity_ids) or "__ground__" in self.entity_ids: + raise ValueError("Collision link and object names must be distinct") + for uid in self.entity_ids: + obj = sim.get_rigid_object(uid) + if not isinstance(obj.cfg.shape, CubeCfg): + raise ValueError( + "PickUp collision validation currently supports cuboid rigid objects only" + ) + scale = obj.get_body_scale().detach().cpu().numpy() + if not np.allclose(scale, scale[:1], atol=0, rtol=0): + raise ValueError( + "Collision shapes must have identical dimensions in every row" + ) + size = np.asarray(obj.cfg.shape.size, dtype=float) * scale[0] + self._box(uid, size) + # SimulationManager's implicit floor, in each arena's local frame. + self._box("__ground__", np.array([1000.0, 1000.0, 100.0])) + self.ground_pose = np.eye(4) + self.ground_pose[2, 3] = -50.001 + self.self_pairs = tuple(combinations(self.links, 2)) + self.world_pairs = tuple( + (link, uid) + for link in self.links + for uid in (*self.entity_ids, "__ground__") + ) + + def _box(self, uid, size): + if size.shape != (3,) or not np.isfinite(size).all() or (size <= 0).any(): + raise ValueError("Cuboid dimensions must be three positive finite values") + self.objects[uid] = self.fcl.CollisionObject(self.fcl.Box(*size)) + self.bounds[uid] = np.stack([-size / 2, size / 2]) + + def link_poses(self, positions, root_pose, extra_links=()): + from embodichain.lab.sim.objects.articulation import Articulation + + local = Articulation.compute_fk( + self.robot, + positions.to(self.robot.device), + link_names=tuple(dict.fromkeys((*self.links, *extra_links))), + qpos_joint_names=tuple(self.robot.joint_names), + to_dict=True, + ) + root = root_pose.detach().cpu().double().numpy() + return { + name: root @ local[name].get_matrix().detach().cpu().double().numpy() + for name in dict.fromkeys((*self.links, *extra_links)) + } + + def collisions(self, link_poses, entity_poses, target_id, allowed): + poses = {**link_poses, **entity_poses, "__ground__": self.ground_pose} + aabbs = {} + for name, pose in poses.items(): + if name not in self.objects: + raise ValueError(f"Missing collision shape: {name}") + self.objects[name].setTransform( + self.fcl.Transform(pose[:3, :3], pose[:3, 3]) + ) + low, high = self.bounds[name] + center = pose[:3, :3] @ ((low + high) / 2) + pose[:3, 3] + radius = np.abs(pose[:3, :3]) @ ((high - low) / 2) + aabbs[name] = (center - radius, center + radius) + pairs = ( + *self.self_pairs, + *self.world_pairs, + *( + (target_id, uid) + for uid in (*self.entity_ids, "__ground__") + if uid != target_id + ), + ) + for first, second in pairs: + if allowed(first, second): + continue + al, ah = aabbs[first] + bl, bh = aabbs[second] + if (al > bh).any() or (bl > ah).any(): + continue + result = self.fcl.CollisionResult() + if self.fcl.collide( + self.objects[first], + self.objects[second], + self.fcl.CollisionRequest(), + result, + ): + return f"{first} / {second}" + return None diff --git a/embodichain/lab/trajectory_generation/integrations/atomic.py b/embodichain/lab/trajectory_generation/integrations/atomic.py new file mode 100644 index 000000000..902638fe1 --- /dev/null +++ b/embodichain/lab/trajectory_generation/integrations/atomic.py @@ -0,0 +1,163 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Explicit export of an offline MoveEndEffector → PickUp compilation.""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + TrajectoryPhase, + TrajectoryTemplate, +) + +if TYPE_CHECKING: + from embodichain.lab.sim.atomic_actions.plans import CompiledTrajectory + from embodichain.lab.sim.objects import Robot + +__all__ = ["export_pickup_templates"] + + +def export_pickup_templates( + compiled: CompiledTrajectory, + robot: Robot, + *, + control_dt: float, + source_id: str = "atomic_pickup", + source_revision: str = "v1", + template_id: str = "reference_0", + validator_id: str = "task_success", + hold_steps: int = 30, + control_part: str = "arm", +) -> tuple[TrajectoryTemplate, ...]: + """Export full-joint candidates with protected approach/close/lift/hold phases. + + Accepts exactly MoveEndEffector followed by PickUp. Segment boundaries come + from the atomic plans; they are never inferred from motion. The duplicate + zero-time boundary must be a shared position and becomes one control-period + hold. Mimic geometry is expanded from the robot's declared affine coupling. + Only the initial transit permits spatial residual augmentation. + + This exports an offline atomic plan for independently validated qpos replay. + It does not commit projected symbolic effects or run AtomicActionRuntime's + recovery/feedback state machine. Physical contact and task verification are + required separately before an exported trajectory can become expert data. + + Args: + compiled: Successful full-batch result from AtomicActionEngine.compile. + robot: Robot whose full joint order and mimic model produced the plan. + control_dt: Fixed replay period, matching every positive plan interval. + source_id: Stable source identity. + source_revision: Caller-owned source revision. + template_id: Template identity shared across row-specific cases. + validator_id: Required physical task check identity. + hold_steps: Explicit terminal commands used to measure a stable grasp. + control_part: Arm joints authorized for transit residuals. + + Returns: + One independently owned template for each physical row. + """ + if ( + not math.isfinite(control_dt) + or control_dt <= 0 + or type(hold_steps) is not int + or hold_steps < 2 + ): + raise ValueError("Provide a positive control_dt and at least two hold steps") + if tuple(plan.skill_id for plan in compiled.action_plans) != ( + "move_end_effector", + "pick_up", + ): + raise ValueError("Export requires exactly MoveEndEffector followed by PickUp") + if not bool(compiled.plan_success.all()): + raise ValueError("Cannot export failed atomic plans") + positions = compiled.trajectory.positions.clone() + intervals = compiled.trajectory.dt.clone() + if positions.shape[0] != robot.num_instances or positions.shape[2] != robot.dof: + raise ValueError( + "Compiled trajectory must cover the complete physical robot batch" + ) + if not torch.equal( + compiled.trajectory.env_ids.cpu(), torch.arange(robot.num_instances) + ): + raise ValueError("Compiled row order must be the physical environment order") + transit = compiled.action_waypoint_offset(1) + approach, close, lift = ( + compiled.segment(1, name) for name in ("approach", "close", "lift") + ) + if ( + approach.start != transit + or approach.stop != close.start + or close.stop != lift.start + or lift.stop != positions.shape[1] + ): + raise ValueError("PickUp segments must cover the complete second action") + if not torch.allclose( + positions[:, transit], positions[:, transit - 1], atol=1e-5, rtol=0 + ): + raise ValueError("Atomic action boundary is not a shared position") + if not bool((intervals[:, 0] == 0).all() and (intervals[:, transit] == 0).all()): + raise ValueError("Expected explicit zero-time action anchors") + intervals[:, transit] = control_dt + if not torch.allclose( + intervals[:, 1:].double(), + torch.full_like(intervals[:, 1:].double(), control_dt), + atol=1e-8, + rtol=0, + ): + raise ValueError("Atomic timing must match the replay control clock") + if set(robot.mimic_ids) & set(robot.mimic_parents): + raise ValueError("PickUp export requires direct active-to-mimic coupling") + for child, parent, multiplier, offset in zip( + robot.mimic_ids, + robot.mimic_parents, + robot.mimic_multipliers, + robot.mimic_offsets, + ): + positions[..., child] = positions[..., parent] * multiplier + offset + length = positions.shape[1] + positions = torch.cat( + (positions, positions[:, -1:].repeat(1, hold_steps, 1)), dim=1 + ) + intervals = torch.cat( + (intervals, intervals.new_full((robot.num_instances, hold_steps), control_dt)), + dim=1, + ) + phases = ( + TrajectoryPhase("transit", 0, transit, allowed_operators=("joint_residual",)), + TrajectoryPhase("approach", transit, close.start, kind="contact"), + TrajectoryPhase("close", close.start, lift.start, kind="contact"), + TrajectoryPhase("lift", lift.start, length, kind="contact"), + TrajectoryPhase("hold", length, length + hold_steps, kind="hold"), + ) + return tuple( + TrajectoryTemplate( + source_id, + source_revision, + template_id, + tuple(robot.joint_names), + positions[row], + intervals[row], + phases, + allowed_operators=("joint_residual",), + validator_id=validator_id, + controlled_joint_indices=tuple(robot.get_joint_ids(control_part)), + ) + for row in range(robot.num_instances) + ) diff --git a/embodichain/lab/trajectory_generation/integrations/contact.py b/embodichain/lab/trajectory_generation/integrations/contact.py new file mode 100644 index 000000000..358a698a3 --- /dev/null +++ b/embodichain/lab/trajectory_generation/integrations/contact.py @@ -0,0 +1,661 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Contact-aware, full-state validation for one fixed-scene cuboid PickUp.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass, field +import math +from typing import TYPE_CHECKING + +import numpy as np +import torch + +from embodichain.utils import configclass +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CandidateTrajectoryBatch, + ExpertEpisode, + MotionSnapshot, + ValidationCheck, + ValidationResult, +) +from ._collision import _FullStateCollisionWorld + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.objects import Robot + +__all__ = ["PickUpContactProfile", "PickUpMotionValidator"] + + +def _result(name, passed, detail="", **metrics): + return ValidationResult( + (ValidationCheck(name, "passed" if passed else "failed", detail, metrics),) + ) + + +@configclass +class PickUpContactProfile: + """Explicit permissions and measured-grasp tolerances for a single PickUp. + + Finger/object contacts are allowed in the declared final approach region + and from ``close`` onwards. Object/support contact is allowed through + initial lift-off; robot/support contact + is restricted to the declared fixed mounting links. Self pairs within two + URDF kinematic hops follow the existing cuRobo structural exclusion policy. + All other robot/world, object/world and nonexcluded self collisions reject. + """ + + object_id: str = "cube" + support_id: str = "bench" + finger_links: tuple[str, ...] = ("gripper_finger1_link_1", "gripper_finger2_link_1") + mounting_links: tuple[str, ...] = ("arm_base_link",) + tcp_link: str = "ee_link" + tcp_offset: tuple[tuple[float, ...], ...] = ( + (1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.15), + (0.0, 0.0, 0.0, 1.0), + ) + min_lift: float = 0.12 + max_relative_translation: float = 0.01 + max_relative_rotation: float = 0.15 + max_tcp_distance: float = 0.06 + approach_contact_distance: float = 0.025 + max_penetration: float = 0.002 + min_contact_fraction: float = 0.95 + min_hold_seconds: float = 1.0 + measured_joint_tolerance: float = 1e-4 + max_joint_step: float = 0.02 + max_validation_samples: int = 4096 + max_contacts_per_step: int = 4096 + + def __post_init__(self) -> None: + for name in ("object_id", "support_id", "tcp_link"): + if not isinstance(getattr(self, name), str) or not getattr(self, name): + raise ValueError(f"{name} must be a nonempty name") + if self.object_id == self.support_id: + raise ValueError("PickUp target and support must differ") + for name in ("finger_links", "mounting_links"): + values = getattr(self, name) + if ( + isinstance(values, str) + or any(not isinstance(v, str) or not v for v in values) + or len(set(values)) != len(values) + ): + raise ValueError(f"{name} must contain unique names") + setattr(self, name, tuple(values)) + if len(self.finger_links) != 2: + raise ValueError( + "This PickUp profile requires two distinct gripping fingers" + ) + from embodichain.lab.sim.motion.trajectory_augmentation.contracts import _pose + + _pose(torch.tensor(self.tcp_offset, dtype=torch.float64), "tcp_offset") + for name in ( + "min_lift", + "max_relative_translation", + "max_relative_rotation", + "max_tcp_distance", + "approach_contact_distance", + "max_penetration", + "min_contact_fraction", + "min_hold_seconds", + "measured_joint_tolerance", + "max_joint_step", + ): + value = getattr(self, name) + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value <= 0 + ): + raise ValueError(f"{name} must be finite and positive") + if self.min_contact_fraction > 1: + raise ValueError("min_contact_fraction must be <= 1") + for name in ("max_validation_samples", "max_contacts_per_step"): + if type(getattr(self, name)) is not int or getattr(self, name) < 2: + raise ValueError(f"{name} must be an integer >= 2") + + +@dataclass +class _Evidence: + steps: int = 0 + hold_steps: int = 0 + contact_steps: list[int] = field(default_factory=lambda: [0, 0]) + first_relative: np.ndarray | None = None + minimum_lift: float = math.inf + max_translation: float = 0.0 + max_rotation: float = 0.0 + max_distance: float = 0.0 + failure: str | None = None + + +class PickUpMotionValidator: + """Validate compiled PickUp plans and actual CPU-physics contact evidence. + + This bounded integration accepts an unscaled, fixed-base URDF robot and + cuboid rigid objects. Planned checks use conservative convex hulls of URDF + *collision* shapes, full joint FK including mimic geometry, a TCP-relative + lifted object, and densified joint samples. Actual checks use measured full + joints and object poses; native contacts are also checked after every physics + substep. No continuous-collision guarantee is claimed. + + The same instance must be attached to QposRolloutExecutor and GenerationRunner. + Only the target object and physically observed mimic joints may move outside + the active controller layout. Gym is deliberately unsupported until its + execution loop exposes equivalent physical-substep evidence. + + Args: + sim: Manually stepped CPU-physics simulation owning every row. + robot: The sole fixed-base robot, with initialized full-articulation FK. + profile: Explicit collision permissions and measured task thresholds. + """ + + def __init__( + self, + sim: SimulationManager, + robot: Robot, + *, + profile: PickUpContactProfile | None = None, + ) -> None: + if torch.device(sim.sim_config.sim_device).type != "cpu": + raise ValueError("PickUp contact evidence currently requires CPU physics") + self.sim, self.robot = sim, robot + self.profile = (profile or PickUpContactProfile()).copy() + self.profile.__post_init__() + self.world = _FullStateCollisionWorld(sim, robot) + p = self.profile + if {p.object_id, p.support_id} - set(self.world.entity_ids): + raise ValueError( + "PickUp object and support must be registered collision entities" + ) + if ( + set((*p.finger_links, *p.mounting_links)) - set(self.world.links) + or p.tcp_link not in robot.link_names + ): + raise ValueError( + "PickUp link permissions must resolve to actual collision links" + ) + self.tcp_offset = torch.tensor( + p.tcp_offset, dtype=torch.float32, device=robot.device + ) + if set(p.mounting_links) - self.world.fixed_links: + raise ValueError( + "Only links fixed to the robot root may contact its mounting support" + ) + self.mimic_ids = tuple(robot.mimic_ids) + self.dynamic_entity_ids = (p.object_id,) + import dexsim + + self._physics = dexsim.default_world().get_physics_scene() + self._physics.enable_contact_data_update_on_cpu(True) + self._users = {} + for link in robot.link_names: + self._register_users(link, robot.get_user_ids(link)) + for uid in self.world.entity_ids: + self._register_users(uid, sim.get_rigid_object(uid).get_user_ids()) + self._rows = () + self._snapshots = () + self._evidence = [] + + @property + def collision_world_entity_ids(self) -> tuple[str, ...]: + """Physical rigid UIDs represented by this complete collision model.""" + return self.world.entity_ids + + def _register_users(self, name, values): + ids = values.detach().cpu().reshape(-1).tolist() + if len(ids) != self.robot.num_instances: + raise ValueError("Contact user IDs must map every physical row") + for row, uid in enumerate(ids): + if uid in self._users and self._users[uid] != (row, name): + raise ValueError("Ambiguous native contact body ID") + self._users[uid] = (row, name) + + def _phases(self, phases, length): + expected = ("transit", "approach", "close", "lift", "hold") + if tuple(phase.phase_id for phase in phases) != expected: + raise ValueError( + "PickUp requires explicit transit/approach/close/lift/hold phases" + ) + end = 0 + kinds = ("free", "contact", "contact", "contact", "hold") + for phase, kind in zip(phases, kinds): + if phase.start_index != end or phase.kind != kind: + raise ValueError("PickUp phases must cover the trajectory in order") + if phase.phase_id != "transit" and phase.allowed_operators: + raise ValueError("PickUp contact and approach phases must be protected") + end = phase.stop_index + if end != length: + raise ValueError("PickUp phases must cover every trajectory sample") + return phases + + def _allowed(self, first, second, phase, lift, grasp_distance=math.inf): + p = self.profile + pair = frozenset((first, second)) + if pair in self.world.adjacent_pairs: + return True + if p.support_id in pair and any(link in pair for link in p.mounting_links): + return True + if p.object_id in pair: + other = second if first == p.object_id else first + if other in p.finger_links and ( + phase in ("close", "lift", "hold") + or ( + phase == "approach" + and grasp_distance <= p.approach_contact_distance + ) + ): + return True + if other == p.support_id and ( + phase in ("transit", "approach", "close") + or (phase == "lift" and lift <= 0.005) + ): + return True + return False + + def _validate_path(self, qpos, phases, snapshot, row, *, measured_objects=None): + p = self.profile + qpos = qpos.detach().cpu() + self._phases(phases, len(qpos)) + if not torch.allclose( + qpos[0], snapshot.joint_positions.cpu().to(qpos), atol=1e-6, rtol=0 + ): + raise ValueError("PickUp path must begin at the captured initial joints") + limits = self.robot.get_qpos_limits()[row].detach().cpu().to(qpos) + tolerance = 0.0 if measured_objects is None else p.measured_joint_tolerance + if bool( + ( + (qpos < limits[:, 0] - tolerance) | (qpos > limits[:, 1] + tolerance) + ).any() + ): + return _result( + "path_collision", + False, + "Full-joint position limits exceeded", + maximum_joint_limit_violation=float( + torch.maximum(limits[:, 0] - qpos, qpos - limits[:, 1]) + .clamp(min=0) + .max() + ), + ) + if set(snapshot.entity_poses) != set(self.world.entity_ids): + raise ValueError("PickUp snapshots must cover the complete collision world") + if measured_objects is None: + # Planned geometry must model passive mimic motion, even though the + # drive only receives targets for active joints. + for child, parent, multiplier, offset in zip( + self.robot.mimic_ids, + self.robot.mimic_parents, + self.robot.mimic_multipliers, + self.robot.mimic_offsets, + ): + if not torch.allclose( + qpos[:, child], + qpos[:, parent] * multiplier + offset, + atol=1e-6, + rtol=0, + ): + raise ValueError( + "Planned mimic coordinates disagree with the robot coupling" + ) + increments = qpos.diff(dim=0).abs().amax(dim=1) + counts = torch.clamp( + torch.ceil(increments / p.max_joint_step).long(), min=1 + ) + if int(counts.sum()) + 1 > p.max_validation_samples: + raise ValueError("PickUp collision sampling exceeds the bounded budget") + samples, source_indices = [qpos[:1]], [0] + for index, count in enumerate(counts.tolist(), 1): + alpha = torch.arange(1, count + 1, dtype=qpos.dtype)[:, None] / count + samples.append(torch.lerp(qpos[index - 1], qpos[index], alpha)) + source_indices.extend([index] * count) + values = torch.cat(samples) + else: + values = qpos + source_indices = list(range(len(qpos))) + if len(values) > p.max_validation_samples or measured_objects.shape != ( + len(values), + 4, + 4, + ): + raise ValueError( + "Actual collision evidence has incompatible shape or exceeds its budget" + ) + poses = self.world.link_poses(values, snapshot.root_pose, (p.tcp_link,)) + tcp = poses[p.tcp_link] @ self.tcp_offset.cpu().double().numpy() + lift_start = phases[3].start_index + anchor = next( + i for i, index in enumerate(source_indices) if index >= lift_start + ) + initial_object = snapshot.entity_poses[p.object_id].cpu().double().numpy() + relative = np.linalg.inv(tcp[anchor]) @ initial_object + for sample, index in enumerate(source_indices): + phase = next( + phase.phase_id + for phase in phases + if phase.start_index <= index < phase.stop_index + ) + objects = { + uid: pose.cpu().double().numpy() + for uid, pose in snapshot.entity_poses.items() + } + if measured_objects is not None: + objects[p.object_id] = measured_objects[sample].cpu().double().numpy() + elif index >= lift_start: + objects[p.object_id] = tcp[sample] @ relative + height = objects[p.object_id][2, 3] - initial_object[2, 3] + collision = self.world.collisions( + {link: poses[link][sample] for link in self.world.links}, + objects, + p.object_id, + lambda a, b: self._allowed( + a, + b, + phase, + height, + float( + np.linalg.norm(tcp[sample, :3, 3] - objects[p.object_id][:3, 3]) + ), + ), + ) + if collision: + return _result( + "path_collision", + False, + f"Forbidden collision at sample {index} ({phase}): {collision}", + checked_samples=sample + 1, + ) + return _result( + "path_collision", + True, + "Conservative full-joint URDF hulls and object/world samples", + checked_samples=len(values), + ) + + def validate_qpos( + self, batch: CandidateTrajectoryBatch, snapshots: Sequence[MotionSnapshot] + ) -> tuple[ValidationResult, ...]: + """Check planned full-state paths, phase permissions and lifted geometry. + + Args: + batch: Full-joint candidates with explicit physical source rows. + snapshots: Initial states in physical environment order. + + Returns: + One path validation result per candidate, in candidate order. + """ + if ( + len(snapshots) != self.robot.num_instances + or batch.source_row_indices is None + or batch.joint_names != tuple(self.robot.joint_names) + ): + raise ValueError( + "PickUp validation requires full ordered joints and real source rows" + ) + results = [] + for index, row in enumerate(batch.source_row_indices.tolist()): + if row >= len(snapshots): + raise ValueError("Source row is outside the physical batch") + snapshot = snapshots[row] + if snapshot.joint_names != batch.joint_names: + raise ValueError( + "PickUp snapshots must use the full ordered robot joints" + ) + identity = batch.identities[index] + if (identity.scene_case_id, identity.initial_state_id) != ( + snapshot.scene_case.scene_case_id, + snapshot.scene_case.initial_state_id, + ): + raise ValueError( + "PickUp candidate identity does not match its snapshot" + ) + length = int(batch.valid_length[index]) + results.append( + self._validate_path( + batch.positions[index, :length], batch.phases[index], snapshot, row + ) + ) + return tuple(results) + + def validate_episode( + self, episode: ExpertEpisode, snapshot: MotionSnapshot, *, row: int + ) -> ValidationResult: + """Recheck measured finger geometry and measured target poses after rollout. + + Args: + episode: Frozen observations containing full joints and object poses. + snapshot: Initial scene state belonging to the episode's case. + row: Physical environment row used during execution. + + Returns: + The measured path's joint-limit and collision validation result. + """ + return self._validate_path( + episode.observations["joint_positions"], + episode.phases, + snapshot, + row, + measured_objects=episode.observations["object_pose"], + ) + + def observations(self) -> dict[str, torch.Tensor]: + """Return current physical object and TCP poses in each local arena. + + Returns: + Object and TCP transforms, each shaped ``(num_envs, 4, 4)``. + """ + return { + "object_pose": self.sim.get_rigid_object( + self.profile.object_id + ).get_local_pose(to_matrix=True), + "tcp_pose": self.robot.get_link_pose(self.profile.tcp_link, to_matrix=True) + @ self.tcp_offset, + } + + def begin_rollout( + self, + candidates: Sequence[CandidateTrajectoryBatch | None], + snapshots: Sequence[MotionSnapshot], + ) -> None: + """Reset bounded per-row evidence before the first physical command. + + Args: + candidates: One candidate per physical row, or ``None`` for idle rows. + snapshots: Initial states in the same physical row order. + """ + if len(candidates) != self.robot.num_instances or len(snapshots) != len( + candidates + ): + raise ValueError( + "Contact monitoring must cover the complete physical batch" + ) + for candidate in candidates: + if candidate is not None: + self._phases(candidate.phases[0], int(candidate.valid_length[0])) + self._rows, self._snapshots = tuple(candidates), tuple(snapshots) + self._evidence = [_Evidence() for _ in candidates] + + def observe_substep( + self, sample_index: int, active: Sequence[bool], *, physics_dt: float + ) -> None: + """Accumulate native contacts and held-object stability after one physics step. + + The CPU contact buffer is read before the next physics update. Unknown + bodies, cross-row contacts, excessive penetration and buffer budget + overflow cannot produce accepted evidence. No contact-history arrays grow. + + Args: + sample_index: Command sample whose physics substep just completed. + active: Whether each physical row is still executing its candidate. + physics_dt: Duration of the completed physics substep in seconds. + """ + p = self.profile + data, users = self._physics.get_cpu_contact_buffer() + data, users = np.asarray(data), np.asarray(users) + if ( + data.ndim != 2 + or data.shape[1] != 11 + or users.shape != (len(data), 2) + or not np.isfinite(data).all() + or len(data) > p.max_contacts_per_step + ): + raise ValueError( + "Native contact evidence is malformed or exceeds its bounded budget" + ) + observation = { + key: value.detach().cpu().double().numpy() + for key, value in self.observations().items() + } + finger_contacts = [set() for _ in active] + phases = [] + for row, candidate in enumerate(self._rows): + phases.append( + None + if candidate is None or not active[row] + else next( + phase.phase_id + for phase in candidate.phases[0] + if phase.start_index <= sample_index < phase.stop_index + ) + ) + for contact, pair in zip(data, users): + first, second = self._users.get(int(pair[0])), self._users.get(int(pair[1])) + if first is None and second is None: + continue + known = first or second + row = known[0] + if first is not None and second is not None and first[0] != second[0]: + for owner in (first[0], second[0]): + if active[owner]: + self._evidence[owner].failure = ( + self._evidence[owner].failure or "Cross-row native contact" + ) + continue + if not active[row]: + continue + evidence = self._evidence[row] + if first is None or second is None: + evidence.failure = ( + evidence.failure or "Unmapped body or cross-row native contact" + ) + continue + a, b = first[1], second[1] + height = observation["object_pose"][row, 2, 3] - float( + self._snapshots[row].entity_poses[p.object_id][2, 3] + ) + distance = float( + np.linalg.norm( + observation["tcp_pose"][row, :3, 3] + - observation["object_pose"][row, :3, 3] + ) + ) + if not self._allowed(a, b, phases[row], height, distance): + evidence.failure = ( + evidence.failure + or f"Forbidden physical contact during {phases[row]}: {a} / {b}" + ) + if contact[10] < -p.max_penetration: + evidence.failure = ( + evidence.failure + or f"Physical penetration exceeds tolerance: {a} / {b}" + ) + if contact[9] > 1e-7 and p.object_id in (a, b): + other = b if a == p.object_id else a + if other in p.finger_links: + finger_contacts[row].add(other) + for row, phase in enumerate(phases): + if phase is None: + continue + evidence = self._evidence[row] + evidence.steps += 1 + if phase not in ("lift", "hold"): + continue + if phase == "hold": + evidence.hold_steps += 1 + for index, link in enumerate(p.finger_links): + evidence.contact_steps[index] += link in finger_contacts[row] + obj, tcp = observation["object_pose"][row], observation["tcp_pose"][row] + relative = np.linalg.inv(tcp) @ obj + if evidence.first_relative is None: + evidence.first_relative = relative.copy() + delta = np.linalg.inv(evidence.first_relative) @ relative + if phase == "hold": + evidence.minimum_lift = min( + evidence.minimum_lift, + obj[2, 3] + - float(self._snapshots[row].entity_poses[p.object_id][2, 3]), + ) + evidence.max_translation = max( + evidence.max_translation, float(np.linalg.norm(delta[:3, 3])) + ) + angle = math.acos(float(np.clip((np.trace(delta[:3, :3]) - 1) / 2, -1, 1))) + evidence.max_rotation = max(evidence.max_rotation, angle) + evidence.max_distance = max( + evidence.max_distance, float(np.linalg.norm(relative[:3, 3])) + ) + self._physics_dt = physics_dt + + def rollout_validation(self, row: int) -> ValidationResult: + """Freeze contact coverage and stable-grasp gates for one completed row. + + Args: + row: Physical environment row monitored by this validator. + + Returns: + Mandatory native-contact and held-object stability checks. + """ + e, p = self._evidence[row], self.profile + duration = e.hold_steps * getattr(self, "_physics_dt", 0.0) + fractions = [value / max(1, e.hold_steps) for value in e.contact_steps] + lift = float(e.minimum_lift) if math.isfinite(e.minimum_lift) else 0.0 + held = ( + duration >= p.min_hold_seconds + and lift >= p.min_lift + and e.max_translation <= p.max_relative_translation + and e.max_rotation <= p.max_relative_rotation + and e.max_distance <= p.max_tcp_distance + and min(fractions) >= p.min_contact_fraction + ) + return ValidationResult( + ( + ValidationCheck( + "physical_contacts", + "passed" if e.steps and e.failure is None else "failed", + e.failure + or "Native contacts checked after every CPU physics substep", + {"physics_samples": e.steps}, + ), + ValidationCheck( + "held_object_stability", + "passed" if held else "failed", + "Both fingers must maintain physical contact with a lifted, stable object", + { + "hold_seconds": duration, + "minimum_lift_m": lift, + "max_relative_translation_m": e.max_translation, + "max_relative_rotation_rad": e.max_rotation, + "max_tcp_distance_m": e.max_distance, + "finger_0_contact_fraction": fractions[0], + "finger_1_contact_fraction": fractions[1], + }, + ), + ) + ) diff --git a/embodichain/lab/trajectory_generation/integrations/planning.py b/embodichain/lab/trajectory_generation/integrations/planning.py new file mode 100644 index 000000000..47c2c379a --- /dev/null +++ b/embodichain/lab/trajectory_generation/integrations/planning.py @@ -0,0 +1,614 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Bounded environment-row planning for explicitly unloaded free trajectories.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +import math +from typing import TYPE_CHECKING + +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CandidateIdentity, + CandidateTrajectoryBatch, + MotionSnapshot, + TrajectoryPhase, + TrajectoryTemplate, + ValidationCheck, + ValidationResult, +) + +if TYPE_CHECKING: + from embodichain.lab.sim.motion.planners import MotionGenerator + +__all__ = ["EEFPath", "EnvRowMotionPlanner"] + + +@dataclass(frozen=True) +class EEFPath: + """Explicit TCP samples in the source row's local arena frame. + + The first pose is the initial TCP anchor. ``dt`` contains arrival intervals, + starting at zero. Optional solved joint samples preserve a chosen IK branch + exactly and are checked against the poses through FK, without another IK. + Paths have no attachment or contact-permission representation. + """ + + identity: CandidateIdentity + source_row_index: int + poses: torch.Tensor + dt: torch.Tensor + phases: tuple[TrajectoryPhase, ...] + solved_joint_targets: torch.Tensor | None = None + + def __post_init__(self) -> None: + if not isinstance(self.identity, CandidateIdentity): + raise ValueError("identity must be CandidateIdentity") + if type(self.source_row_index) is not int or self.source_row_index < 0: + raise ValueError("source_row_index must be a nonnegative integer") + poses = self.poses + if ( + not isinstance(poses, torch.Tensor) + or not poses.is_floating_point() + or poses.ndim != 3 + or poses.shape[1:] != (4, 4) + or not poses.shape[0] + or not bool(torch.isfinite(poses).all()) + ): + raise ValueError("poses must be finite floating shape (N >= 1, 4, 4)") + rotation = poses[:, :3, :3].double() + identity = torch.eye(3, device=poses.device, dtype=torch.float64) + if ( + not torch.allclose( + rotation.transpose(1, 2) @ rotation, + identity.expand_as(rotation), + atol=1e-5, + rtol=0, + ) + or not torch.allclose( + torch.linalg.det(rotation), + torch.ones(len(poses), device=poses.device, dtype=torch.float64), + atol=1e-5, + rtol=0, + ) + or not torch.allclose( + poses[:, 3], + poses.new_tensor([0, 0, 0, 1]).expand(len(poses), 4), + atol=1e-6, + rtol=0, + ) + ): + raise ValueError("poses must contain proper SE(3) transforms") + reference = TrajectoryTemplate( + self.identity.source_id, + self.identity.source_revision, + self.identity.template_id, + ("validation_joint",), + poses.new_zeros((len(poses), 1)), + self.dt, + self.phases, + ) + solved = self.solved_joint_targets + if solved is not None: + if ( + not isinstance(solved, torch.Tensor) + or not solved.is_floating_point() + or solved.ndim != 2 + or solved.shape[0] != len(poses) + or not solved.shape[1] + or not bool(torch.isfinite(solved).all()) + ): + raise ValueError( + "solved_joint_targets must be finite floating shape (N, D)" + ) + solved = solved.detach().clone() + object.__setattr__(self, "poses", poses.detach().clone()) + object.__setattr__(self, "dt", reference.dt) + object.__setattr__(self, "phases", reference.phases) + object.__setattr__(self, "solved_joint_targets", solved) + + +def _check(status: str, detail: str, **metrics: float) -> ValidationResult: + return ValidationResult( + (ValidationCheck("path_collision", status, detail, metrics),) + ) + + +class EnvRowMotionPlanner: + """Convert EEF paths and check qpos without changing simulator batch size. + + Only fully annotated free motion with no held object is supported. Every + uncontrolled joint must remain at the snapshot and robot configuration's + initial value, matching the current planner's locked-joint model. Callers + keep the live root poses and collision model synchronized with these same + snapshots. Scene entities use canonical IDs in the backend collision world. + + Collision validation densifies joint segments, including phase boundaries, + then checks those samples through ``MotionGenerator``. This is a bounded + sampled approximation, not continuous collision detection. It neither steps + physics nor verifies task success or controller velocity/acceleration limits. + + Args: + motion_generator: Existing generator attached to the real batch robot. + control_part: Ordered control part represented by the backend model. + held_object_ids: Explicit held-object declaration; nonempty is unsupported. + max_joint_step: Maximum absolute joint increment between checked samples. + max_validation_samples: Maximum samples per candidate, including anchors. + fk_position_tolerance: Maximum TCP translation residual in meters. + fk_rotation_tolerance: Maximum TCP rotation residual in radians. + """ + + def __init__( + self, + motion_generator: MotionGenerator, + *, + control_part: str, + held_object_ids: Sequence[str], + max_joint_step: float = 0.02, + max_validation_samples: int = 4096, + fk_position_tolerance: float = 1e-3, + fk_rotation_tolerance: float = 1e-2, + ) -> None: + self.motion_generator = motion_generator + self.robot = motion_generator.robot + if not isinstance(control_part, str) or not control_part: + raise ValueError("control_part must be a nonempty name") + self.control_part = control_part + self.joint_names = tuple(self.robot.joint_names) + self.joint_ids = tuple(self.robot.get_joint_ids(control_part)) + if not self.joint_ids or len(set(self.joint_ids)) != len(self.joint_ids): + raise ValueError("control_part must contain unique joints") + self.fixed_ids = tuple( + i for i in range(len(self.joint_names)) if i not in self.joint_ids + ) + if isinstance(held_object_ids, str) or any( + not isinstance(x, str) or not x for x in held_object_ids + ): + raise ValueError("held_object_ids must be an explicit sequence of names") + self.held_object_ids = tuple(held_object_ids) + for name, value in ( + ("max_joint_step", max_joint_step), + ("fk_position_tolerance", fk_position_tolerance), + ("fk_rotation_tolerance", fk_rotation_tolerance), + ): + if ( + isinstance(value, bool) + or not isinstance(value, (float, int)) + or not math.isfinite(value) + or value <= 0 + ): + raise ValueError(f"{name} must be positive and finite") + setattr(self, name, float(value)) + if type(max_validation_samples) is not int or max_validation_samples < 1: + raise ValueError("max_validation_samples must be a positive integer") + self.max_validation_samples = max_validation_samples + + def _snapshots( + self, snapshots: Sequence[MotionSnapshot] + ) -> tuple[MotionSnapshot, ...]: + values = tuple(snapshots) + if len(values) != self.robot.num_instances or not values: + raise ValueError("snapshots must cover the complete real robot batch") + if any( + not isinstance(value, MotionSnapshot) + or value.joint_names != self.joint_names + for value in values + ): + raise ValueError("snapshots must use the robot's full ordered joint names") + current_roots = self.robot.get_local_pose(to_matrix=True) + for row, snapshot in enumerate(values): + if not torch.allclose( + current_roots[row].to(snapshot.root_pose), + snapshot.root_pose, + atol=1e-6, + rtol=0, + ): + raise ValueError("snapshot root pose differs from its live source row") + return values + + @staticmethod + def _rounds(indices: Sequence[int], candidate_ids: Sequence[int]): + queues: dict[int, list[int]] = {} + for candidate in candidate_ids: + queues.setdefault(indices[candidate], []).append(candidate) + while any(queues.values()): + yield {row: queue.pop(0) for row, queue in queues.items() if queue} + + def _unsupported( + self, phases: tuple[TrajectoryPhase, ...], length: int + ) -> str | None: + if self.held_object_ids: + return "held-object sweep geometry is unavailable" + end = 0 + for phase in phases: + if phase.kind != "free" or phase.start_index != end: + return ( + "only fully annotated free motion has supported collision semantics" + ) + end = phase.stop_index + return ( + None if end == length else "every sample must have an explicit free phase" + ) + + def _obstacles(self, snapshots, obstacle_poses): + names = self.motion_generator.dynamic_collision_entity_ids + if obstacle_poses is None and names: + if any( + name not in snapshot.entity_poses + for name in names + for snapshot in snapshots + ): + raise ValueError("snapshots lack required dynamic collision poses") + obstacle_poses = { + name: torch.stack( + [snapshot.entity_poses[name] for snapshot in snapshots] + ) + for name in names + } + if obstacle_poses is None: + return None + if set(obstacle_poses) != set(names): + raise ValueError( + "obstacle poses must match the backend's dynamic entity IDs" + ) + result = {} + for name, pose in obstacle_poses.items(): + if ( + not isinstance(pose, torch.Tensor) + or pose.shape != (len(snapshots), 4, 4) + or not bool(torch.isfinite(pose).all()) + ): + raise ValueError("obstacle poses must cover all real source rows") + result[name] = pose.detach().clone() + return result + + def validate_qpos( + self, + batch: CandidateTrajectoryBatch, + snapshots: Sequence[MotionSnapshot], + *, + obstacle_poses: Mapping[str, torch.Tensor] | None = None, + ) -> tuple[ValidationResult, ...]: + """Check original candidates without replacing their positions or timing. + + Return one result per candidate in input order. Unsupported geometry or + backend capability yields ``unavailable``; collisions and invalid paths + yield ``failed``. No success is inferred from padding or planning alone. + + Measured qpos may be checked after execution while the live base frames + remain fixed. The first sample must match its snapshot within ``1e-6``. + Obstacle poses describe one constant world for the entire path; callers + must independently establish that the world stayed fixed throughout a + measured rollout. Neither initial nor final obstacle poses reconstruct + a moving obstacle's history. + """ + snapshots = self._snapshots(snapshots) + if batch.joint_names != self.joint_names: + raise ValueError( + "candidate joint names must match the complete robot order" + ) + count = len(batch.identities) + if not count: + return () + if batch.source_row_indices is None: + raise ValueError( + "source_row_indices are required for environment-row planning" + ) + indices = batch.source_row_indices.tolist() + if min(indices) < 0 or max(indices) >= len(snapshots): + raise ValueError("source_row_indices exceed the real robot batch") + obstacles = self._obstacles(snapshots, obstacle_poses) + device = self.robot.device + starts = torch.stack([snapshot.joint_positions for snapshot in snapshots]).to( + device + ) + initial_values = self.robot.cfg.init_qpos + initial = torch.as_tensor( + [] if initial_values is None else initial_values, + device=device, + dtype=starts.dtype, + ) + limits = self.robot.get_qpos_limits().to(device) + results = [ + _check("not_run", "candidate has not been checked") for _ in range(count) + ] + sample_counts: dict[int, int] = {} + for candidate, source in enumerate(indices): + identity = batch.identities[candidate] + case = snapshots[source].scene_case + if (identity.scene_case_id, identity.initial_state_id) != ( + case.scene_case_id, + case.initial_state_id, + ): + raise ValueError( + "candidate identity does not match its source snapshot" + ) + length = int(batch.valid_length[candidate]) + q = batch.positions[candidate, :length].to(device) + dt = batch.dt[candidate, :length] + unsupported = self._unsupported(batch.phases[candidate], length) + if unsupported: + results[candidate] = _check("unavailable", unsupported) + continue + if ( + not bool(torch.isfinite(q).all()) + or not bool(torch.isfinite(dt).all()) + or dt[0] != 0 + or bool((dt[1:] <= 0).any()) + ): + results[candidate] = _check( + "failed", "non-finite positions or invalid arrival intervals" + ) + continue + if not torch.allclose(q[0], starts[source].to(q), atol=1e-6, rtol=0): + results[candidate] = _check( + "failed", "trajectory does not start at the source initial state" + ) + continue + if self.fixed_ids and ( + initial.shape != starts.shape[1:] + or not torch.allclose( + q[:, self.fixed_ids], + starts[source, self.fixed_ids].expand(length, -1).to(q), + atol=1e-6, + rtol=0, + ) + or not torch.allclose( + starts[source, self.fixed_ids], + initial[list(self.fixed_ids)], + atol=1e-6, + rtol=0, + ) + ): + results[candidate] = _check( + "unavailable", + "changing or noninitial locked joints require a full-state collision model", + ) + continue + if bool(((q < limits[source, :, 0]) | (q > limits[source, :, 1])).any()): + results[candidate] = _check("failed", "joint limits exceeded") + continue + if not self.motion_generator.supports_joint_trajectory_validation: + results[candidate] = _check( + "unavailable", "backend cannot validate exact joint samples" + ) + continue + increments = torch.ceil( + (q[1:, self.joint_ids].double() - q[:-1, self.joint_ids].double()) + .abs() + .amax(dim=1) + / self.max_joint_step + ).clamp_min(1) + if float(increments.sum()) + 1 > self.max_validation_samples: + results[candidate] = _check( + "unavailable", "collision sampling exceeds max_validation_samples" + ) + continue + sample_counts[candidate] = int(increments.sum()) + 1 + for round_rows in self._rounds(indices, list(sample_counts)): + horizon = max(sample_counts[index] for index in round_rows.values()) + inputs = starts[:, None, self.joint_ids].expand(-1, horizon, -1).clone() + for row, candidate in round_rows.items(): + length = int(batch.valid_length[candidate]) + q = batch.positions[candidate, :length, self.joint_ids].to(device) + increments = ( + torch.ceil( + (q[1:].double() - q[:-1].double()).abs().amax(dim=1) + / self.max_joint_step + ) + .clamp_min(1) + .long() + ) + cursor = 1 + inputs[row, 0] = q[0] + for index, intervals in enumerate(increments.tolist()): + alpha = ( + torch.arange( + 1, intervals + 1, device=device, dtype=torch.float64 + ) + / intervals + ) + samples = q[index].double() + alpha[:, None] * ( + q[index + 1].double() - q[index].double() + ) + inputs[row, cursor : cursor + intervals] = samples.to(inputs) + cursor += intervals + inputs[row, cursor:] = q[-1] + try: + valid = self.motion_generator.validate_joint_trajectory( + inputs, control_part=self.control_part, obstacle_poses=obstacles + ) + except (ValueError, RuntimeError, NotImplementedError) as error: + for candidate in round_rows.values(): + results[candidate] = _check( + "unavailable", f"backend validation failed: {error}" + ) + continue + for row, candidate in round_rows.items(): + mask = valid[row, : sample_counts[candidate]] + results[candidate] = _check( + "passed" if bool(mask.all()) else "failed", + "sampled joint bounds, self and environment collision checks; no continuous guarantee", + samples=float(len(mask)), + max_joint_step=self.max_joint_step, + ) + return tuple(results) + + def plan_eef( + self, + paths: Sequence[EEFPath], + snapshots: Sequence[MotionSnapshot], + *, + obstacle_poses: Mapping[str, torch.Tensor] | None = None, + ) -> tuple[CandidateTrajectoryBatch, tuple[ValidationResult, ...]]: + """Solve exact EEF samples, preserve supplied branches, and check qpos. + + Returned rows retain input order, phases, intervals, and source indices. + Failed rows hold their full initial state and must be rejected using the + paired validation result. IK always uses the real environment batch. + """ + snapshots = self._snapshots(snapshots) + paths = tuple(paths) + if any(not isinstance(path, EEFPath) for path in paths): + raise TypeError("paths must contain EEFPath values") + if any(path.source_row_index >= len(snapshots) for path in paths): + raise ValueError("EEF source row exceeds the real robot batch") + device = self.robot.device + starts = torch.stack([snapshot.joint_positions for snapshot in snapshots]).to( + device + ) + indices = [path.source_row_index for path in paths] + horizon = max((len(path.poses) for path in paths), default=0) + if horizon > self.max_validation_samples: + raise ValueError("EEF input horizon exceeds max_validation_samples") + for path in paths: + if path.solved_joint_targets is not None: + starts = starts.to( + dtype=torch.promote_types( + starts.dtype, path.solved_joint_targets.dtype + ) + ) + positions = starts.new_empty((len(paths), horizon, len(self.joint_names))) + timing_dtype = paths[0].dt.dtype if paths else starts.dtype + for path in paths[1:]: + timing_dtype = torch.promote_types(timing_dtype, path.dt.dtype) + dt = torch.zeros((len(paths), horizon), device=device, dtype=timing_dtype) + failures: dict[int, ValidationResult] = {} + for index, path in enumerate(paths): + positions[index] = starts[path.source_row_index] + dt[index, : len(path.dt)] = path.dt.to(dt) + unsupported = self._unsupported(path.phases, len(path.poses)) + if unsupported: + failures[index] = _check("unavailable", unsupported) + if ( + path.solved_joint_targets is not None + and path.solved_joint_targets.shape[1] != len(self.joint_ids) + ): + raise ValueError( + "solved_joint_targets must match the ordered control part" + ) + for round_rows in self._rounds( + indices, [i for i in range(len(paths)) if i not in failures] + ): + seeds = starts[:, self.joint_ids].clone() + active = dict(round_rows) + for sample in range( + max(len(paths[index].poses) for index in active.values()) + ): + expected = self.robot.compute_fk( + qpos=seeds, name=self.control_part, to_matrix=True + ) + targets = expected.clone() + unsolved = [] + supplied = {} + for row, candidate in active.items(): + path = paths[candidate] + if sample >= len(path.poses): + continue + targets[row] = path.poses[sample].to(targets) + if sample and path.solved_joint_targets is None: + unsolved.append(row) + elif path.solved_joint_targets is not None: + supplied[row] = path.solved_joint_targets[sample].to(seeds) + solved = seeds.clone() + success = torch.ones(len(snapshots), dtype=torch.bool, device=device) + if unsolved: + ik_targets = expected.clone() + ik_targets[unsolved] = targets[unsolved] + try: + mask, ik = self.robot.compute_ik( + pose=ik_targets, name=self.control_part, joint_seed=seeds + ) + mask = torch.as_tensor( + mask, device=device, dtype=torch.bool + ).reshape(-1) + ik = torch.as_tensor(ik, device=device, dtype=seeds.dtype) + if ik.shape == (len(snapshots), 1, len(self.joint_ids)): + ik = ik[:, 0] + if mask.shape != success.shape or ik.shape != seeds.shape: + raise ValueError( + "IK result has an invalid real-batch shape" + ) + mask &= torch.isfinite(ik).all(dim=1) + success[unsolved] = mask[unsolved] + solved[unsolved] = torch.where( + mask[unsolved, None], ik[unsolved], seeds[unsolved] + ) + except (ValueError, RuntimeError, TypeError) as error: + for row in unsolved: + failures[active[row]] = _check( + "failed", f"IK failed: {error}" + ) + success[row] = False + for row, value in supplied.items(): + solved[row] = value + actual = self.robot.compute_fk( + qpos=solved, name=self.control_part, to_matrix=True + ) + translation = torch.linalg.vector_norm( + actual[:, :3, 3] - targets[:, :3, 3], dim=1 + ) + relative = actual[:, :3, :3].transpose(1, 2) @ targets[:, :3, :3] + cosine = ((relative.diagonal(dim1=1, dim2=2).sum(-1) - 1) / 2).clamp( + -1, 1 + ) + angle = torch.acos(cosine) + success &= ( + torch.isfinite(actual).all(dim=2).all(dim=1) + & (translation <= self.fk_position_tolerance) + & (angle <= self.fk_rotation_tolerance) + ) + for row, candidate in tuple(active.items()): + if sample >= len(paths[candidate].poses): + continue + if not success[row]: + failures.setdefault( + candidate, + _check( + "failed", + "IK failed or TCP FK residual exceeded tolerance", + ), + ) + positions[candidate] = starts[row] + del active[row] + else: + positions[candidate, sample, self.joint_ids] = solved[row] + seeds[row] = solved[row] + if not active: + break + for index, path in enumerate(paths): + positions[index, len(path.poses) :] = positions[index, len(path.poses) - 1] + batch = CandidateTrajectoryBatch( + positions, + dt, + torch.tensor( + [len(path.poses) for path in paths], dtype=torch.int64, device=device + ), + tuple(path.identity for path in paths), + self.joint_names, + tuple(path.phases for path in paths), + source_row_indices=torch.tensor(indices, dtype=torch.int64, device=device), + ) + results = list( + self.validate_qpos(batch, snapshots, obstacle_poses=obstacle_poses) + ) + for index, failure in failures.items(): + results[index] = failure + return batch, tuple(results) diff --git a/embodichain/lab/trajectory_generation/integrations/sim.py b/embodichain/lab/trajectory_generation/integrations/sim.py new file mode 100644 index 000000000..0e01f194c --- /dev/null +++ b/embodichain/lab/trajectory_generation/integrations/sim.py @@ -0,0 +1,433 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Strict physical initial states for an exclusively owned simulation batch.""" + +from __future__ import annotations + +import hashlib +import json +import math +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING + +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + ValidationCheck, + ValidationResult, +) + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.objects import Robot + +__all__ = ["SimInitialState", "SimInitialStateAdapter"] + +_ROBOT_FIELDS = ( + "root_pose", + "root_linear_velocity", + "root_angular_velocity", + "qpos", + "qvel", + "target_qpos", + "target_qvel", + "qf", +) +_RIGID_FIELDS = ("pose", "linear_velocity", "angular_velocity") + + +def _copy_fields(values: Mapping[str, torch.Tensor]) -> Mapping[str, torch.Tensor]: + """Own finite floating-point state independently of backend tensor buffers.""" + copied = {} + for key, value in values.items(): + if not isinstance(key, str) or not isinstance(value, torch.Tensor): + raise TypeError("State fields must map names to tensors.") + if not value.is_floating_point() or not bool(torch.isfinite(value).all()): + raise ValueError( + f"State field {key!r} must contain finite floating values." + ) + copied[key] = value.detach().clone() + return MappingProxyType(copied) + + +@dataclass(frozen=True) +class SimInitialState: + """Owned physical state of every row in one fixed-base robot scene. + + Obtain this value from :meth:`SimInitialStateAdapter.capture`. Pose fields + are local homogeneous matrices; joint fields include mimic and gripper + joints in ``joint_names`` order. Nested mappings are read-only and tensors + are cloned on construction. The adapter validates them again before use. + This value contains no contact solver checkpoint, pending forces, manager + state, or physical/visual property snapshot. + + Args: + batch_size: Number of simulation rows. + robot_uid: Registered robot identity. + joint_names: Complete movable-joint order, including mimic joints. + signature: Structural signature supplied by the capturing adapter. + robot: Root, joint, drive-target, and joint-effort state tensors. + rigid_objects: Pose and velocity tensors keyed by every rigid UID. + """ + + batch_size: int + robot_uid: str + joint_names: tuple[str, ...] + signature: str + robot: Mapping[str, torch.Tensor] + rigid_objects: Mapping[str, Mapping[str, torch.Tensor]] + + def __post_init__(self) -> None: + if type(self.batch_size) is not int or self.batch_size < 1: + raise ValueError("batch_size must be a positive integer.") + if not isinstance(self.robot_uid, str) or not self.robot_uid: + raise ValueError("robot_uid must be a nonempty string.") + if not isinstance(self.signature, str) or not self.signature: + raise ValueError("signature must be a nonempty string.") + names = tuple(self.joint_names) + if ( + not names + or any(not isinstance(name, str) or not name for name in names) + or len(set(names)) != len(names) + ): + raise ValueError("joint_names must contain unique nonempty names.") + object.__setattr__(self, "joint_names", names) + object.__setattr__(self, "robot", _copy_fields(self.robot)) + rigid = {} + for uid, fields in self.rigid_objects.items(): + if not isinstance(uid, str) or not uid: + raise ValueError("Rigid object UIDs must be nonempty strings.") + rigid[uid] = _copy_fields(fields) + object.__setattr__(self, "rigid_objects", MappingProxyType(rigid)) + + +class SimInitialStateAdapter: + """Capture, restore, and verify one entire simulation batch. + + Supports one config-declared fixed-base robot and all ordinary rigid + objects. Additional articulations, object groups, deformables, and rigid + constraints are rejected. Asset-owned USD articulation properties are + rejected because ``cfg.fix_base`` is not authoritative for those assets. + The host owns exclusive batch access, deterministic task initialization, + fixed scene properties, settling, observation refresh, and runtime epochs. + Initial states must have no pending external forces or contact constraints. + + .. attention:: + Restoring the robot root calls the existing articulation pose setter, + which can advance the whole world by 1 ms. Root restoration therefore + precedes all remaining writes. No Gym step, reset event, or recording + callback is invoked. The host must settle and verify before execution. + + Args: + sim: Simulation manager whose complete batch is exclusively owned. + robot: The only robot registered in ``sim``. + atol: Absolute element-wise tolerance for physical verification and + zero root/non-dynamic velocities. Relative tolerance is zero. + tolerances_profile_id: Trusted integration ID naming this tolerance policy. + """ + + def __init__( + self, + sim: SimulationManager, + robot: Robot, + *, + atol: float = 1e-5, + tolerances_profile_id: str = "fixed_scene_tolerances", + ) -> None: + if isinstance(atol, bool) or not isinstance(atol, (int, float)): + raise TypeError("atol must be a finite nonnegative number.") + if not math.isfinite(atol) or atol < 0: + raise ValueError("atol must be a finite nonnegative number.") + self.sim = sim + self.robot = robot + self.atol = float(atol) + if ( + not isinstance(tolerances_profile_id, str) + or not tolerances_profile_id.strip() + ): + raise ValueError("tolerances_profile_id must be a nonempty string.") + self.tolerances_profile_id = tolerances_profile_id + self._check_scene() + + def _check_scene(self) -> None: + """Reject unsupported or incompletely represented scenes before writes.""" + sim, robot = self.sim, self.robot + if sim._robots != {robot.uid: robot}: + raise ValueError( + "Initial states require exactly the supplied registered robot." + ) + if not robot.cfg.fix_base or getattr(robot.cfg, "use_usd_properties", False): + raise ValueError("Initial states require a config-owned fixed-base robot.") + for name in ( + "_articulations", + "_rigid_object_groups", + "_soft_objects", + "_cloth_objects", + "_constraints", + ): + if getattr(sim, name): + raise ValueError(f"Initial state restoration does not support {name}.") + if robot.num_instances != sim.num_envs: + raise ValueError("The robot must cover the entire simulation batch.") + names = tuple(robot.joint_names) + if len(names) != robot.dof or len(set(names)) != len(names) or not names: + raise ValueError( + "Robot joint_names must identify every movable joint once." + ) + for uid, obj in sim._rigid_objects.items(): + if uid != obj.uid or obj.num_instances != sim.num_envs: + raise ValueError( + f"Rigid object {uid!r} does not match its full-batch registration." + ) + + def signature(self) -> str: + """Return a stable topology and control-part signature of the batch. + + Includes row count, robot identity, complete joint order, named control + parts, and rigid object identities/body modes. Current poses and fixed + physical/visual properties are excluded; the host's preparation profile + must separately identify and verify those fixed case conditions. + """ + self._check_scene() + payload = { + "batch_size": self.sim.num_envs, + "robot_uid": self.robot.uid, + "joint_names": list(self.robot.joint_names), + "control_parts": self.robot.control_parts or {}, + "rigid_objects": [ + (uid, bool(obj.is_static), bool(obj.is_non_dynamic)) + for uid, obj in sorted(self.sim._rigid_objects.items()) + ], + } + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() + + def _read(self) -> SimInitialState: + """Read every supported field without invoking simulation updates.""" + signature = self.signature() + robot = self.robot + root_index = list(robot.body_data.link_names).index(robot.root_link_name) + root_velocity = robot.body_data.body_link_vel[:, root_index] + fields = { + "root_pose": robot.get_local_pose(to_matrix=True), + "root_linear_velocity": root_velocity[:, :3], + "root_angular_velocity": root_velocity[:, 3:], + "qpos": robot.get_qpos(), + "qvel": robot.get_qvel(), + "target_qpos": robot.get_qpos(target=True), + "target_qvel": robot.get_qvel(target=True), + "qf": robot.get_qf(), + } + rigid = {} + for uid, obj in self.sim._rigid_objects.items(): + body_state = obj.body_state + rigid[uid] = { + "pose": obj.get_local_pose(to_matrix=True), + "linear_velocity": body_state[:, 7:10], + "angular_velocity": body_state[:, 10:13], + } + return SimInitialState( + batch_size=self.sim.num_envs, + robot_uid=robot.uid, + joint_names=tuple(robot.joint_names), + signature=signature, + robot=fields, + rigid_objects=rigid, + ) + + def _check_fields( + self, + fields: Mapping[str, torch.Tensor], + shapes: Mapping[str, tuple[int, ...]], + label: str, + ) -> None: + if set(fields) != set(shapes): + raise ValueError( + f"{label} fields do not match the complete initial state schema." + ) + for name, shape in shapes.items(): + value = fields[name] + if ( + not isinstance(value, torch.Tensor) + or value.shape != shape + or not value.is_floating_point() + or not bool(torch.isfinite(value).all()) + ): + raise ValueError( + f"{label}.{name} requires finite floating values of shape {shape}." + ) + if name.endswith("pose"): + pose = value.detach().to(dtype=torch.float64, device="cpu") + rotation = pose[:, :3, :3] + if ( + not torch.allclose( + pose[:, 3], + torch.tensor([0.0, 0.0, 0.0, 1.0], dtype=pose.dtype).expand( + shape[0], -1 + ), + atol=1e-5, + rtol=0, + ) + or not torch.allclose( + rotation.mT @ rotation, + torch.eye(3, dtype=pose.dtype).expand(shape[0], -1, -1), + atol=1e-5, + rtol=0, + ) + or not torch.allclose( + torch.linalg.det(rotation), + torch.ones(shape[0], dtype=pose.dtype), + atol=1e-5, + rtol=0, + ) + ): + raise ValueError(f"{label}.{name} must contain rigid SE(3) poses.") + + def _check_state(self, state: SimInitialState) -> None: + """Validate the full snapshot before any physical mutation.""" + if not isinstance(state, SimInitialState): + raise TypeError("state must be a SimInitialState.") + if ( + state.signature != self.signature() + or state.batch_size != self.sim.num_envs + or state.robot_uid != self.robot.uid + or state.joint_names != tuple(self.robot.joint_names) + or set(state.rigid_objects) != set(self.sim._rigid_objects) + ): + raise ValueError( + "Initial state topology, joint order, or controller signature changed." + ) + batch, dof = state.batch_size, self.robot.dof + shapes = {key: (batch, dof) for key in _ROBOT_FIELDS} + shapes.update( + root_pose=(batch, 4, 4), + root_linear_velocity=(batch, 3), + root_angular_velocity=(batch, 3), + ) + self._check_fields(state.robot, shapes, "robot") + for key in ("root_linear_velocity", "root_angular_velocity"): + if bool((state.robot[key].abs() > self.atol).any()): + raise ValueError( + "A fixed-base initial state must have zero root velocity." + ) + limits = ( + self.robot.get_qpos_limits().detach().to(device="cpu", dtype=torch.float64) + ) + if ( + limits.shape != (batch, dof, 2) + or bool(torch.isnan(limits).any()) + or bool((limits[..., 0] > limits[..., 1]).any()) + ): + raise ValueError("Robot joint limits do not cover the full ordered batch.") + for key in ("qpos", "target_qpos"): + value = state.robot[key].detach().to(device="cpu", dtype=torch.float64) + if bool(((value < limits[..., 0]) | (value > limits[..., 1])).any()): + raise ValueError( + f"robot.{key} is outside joint limits and would be clamped." + ) + for uid, fields in state.rigid_objects.items(): + shapes = { + "pose": (batch, 4, 4), + "linear_velocity": (batch, 3), + "angular_velocity": (batch, 3), + } + self._check_fields(fields, shapes, f"rigid_objects.{uid}") + if self.sim._rigid_objects[uid].is_non_dynamic and any( + bool((fields[key].abs() > self.atol).any()) for key in _RIGID_FIELDS[1:] + ): + raise ValueError( + f"Non-dynamic object {uid!r} must have zero initial velocity." + ) + + def capture(self) -> SimInitialState: + """Own a validated physical snapshot after host preparation and settling.""" + state = self._read() + self._check_state(state) + return state + + def restore(self, state: SimInitialState) -> None: + """Restore all rows after complete preflight, without layout randomization. + + Resets pending rigid-body forces/torques, then restores captured + velocities. Fixed-base root velocities must be zero. This method does + not restore arbitrary contact states or hide write/backend failures. + The caller must invalidate its epoch before entry and settle/verify + after completion; an exception can leave partially restored state. + """ + self._check_state(state) + robot = self.robot + fields = { + key: value.detach().to(device=robot.device, dtype=torch.float32).clone() + for key, value in state.robot.items() + } + robot.set_local_pose(fields["root_pose"]) + robot.set_qpos(fields["qpos"], target=False) + robot.set_qvel(fields["qvel"], target=False) + robot.set_qpos(fields["target_qpos"], target=True) + robot.set_qvel(fields["target_qvel"], target=True) + robot.set_qf(fields["qf"]) + for uid, source in state.rigid_objects.items(): + obj = self.sim._rigid_objects[uid] + fields = { + key: value.detach().to(device=obj.device, dtype=torch.float32).clone() + for key, value in source.items() + } + obj.set_local_pose(fields["pose"]) + if not obj.is_non_dynamic: + obj.clear_dynamics() + obj.set_velocity( + lin_vel=fields["linear_velocity"], + ang_vel=fields["angular_velocity"], + ) + + def verify(self, state: SimInitialState) -> ValidationResult: + """Check every physical field after settling, with zero relative tolerance. + + Structural incompatibility or a malformed snapshot raises before + reading physical state. Live nonfinite state and numerical drift + return a failed ``initial_state`` check, never an accepted result. + """ + self._check_state(state) + try: + actual = self._read() + except ValueError as error: + return ValidationResult( + (ValidationCheck("initial_state", "failed", str(error)),) + ) + pairs = [("robot", state.robot, actual.robot)] + pairs.extend( + (f"rigid_objects.{uid}", fields, actual.rigid_objects[uid]) + for uid, fields in state.rigid_objects.items() + ) + for label, expected, observed in pairs: + for key, value in expected.items(): + current = observed[key].to(device=value.device, dtype=value.dtype) + if current.shape != value.shape or not torch.allclose( + current, value, atol=self.atol, rtol=0 + ): + return ValidationResult( + ( + ValidationCheck( + "initial_state", + "failed", + f"{label}.{key} differs from the specified initial state.", + ), + ) + ) + return ValidationResult((ValidationCheck("initial_state", "passed"),)) diff --git a/embodichain/lab/trajectory_generation/runner.py b/embodichain/lab/trajectory_generation/runner.py new file mode 100644 index 000000000..6b5e5e7a6 --- /dev/null +++ b/embodichain/lab/trajectory_generation/runner.py @@ -0,0 +1,743 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Synchronous fixed-scene qpos collection through qualified motion/host ports.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import MISSING, fields, is_dataclass, replace +import hashlib +import json +import math +import os +from pathlib import Path +import tempfile +import time +from typing import TYPE_CHECKING + +import torch + +from embodichain import __version__ +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CandidateTrajectoryBatch, + CommitReceipt, + ExpertEpisode, + GenerationSession, + SceneCase, + TrajectoryGenerationJobCfg, + TrajectoryTemplate, + ValidationCheck, + ValidationResult, + joint_residual, + retime, + validate_motion_limits, +) +from embodichain.utils import configclass + +if TYPE_CHECKING: + from .execution import QposRolloutExecutor + from .initial_state import FixedSceneHost + from .integrations.planning import EnvRowMotionPlanner + from .sinks import LeRobotEpisodeSink + +from .integrations.contact import PickUpMotionValidator + +__all__ = ["MotionLimitsProfile", "GenerationRunner"] + + +@configclass +class MotionLimitsProfile: + """Trusted full-joint velocity and acceleration limits. + + Args: + velocity_limits: Positive finite speed limits in joint units per second. + acceleration_limits: Positive finite acceleration limits in joint units + per second squared, in the same full-joint order. + profile_id: Registered policy ID used by the generation job. + """ + + velocity_limits: torch.Tensor = MISSING + acceleration_limits: torch.Tensor = MISSING + profile_id: str = "robot_execution_limits" + + def __post_init__(self) -> None: + if not isinstance(self.profile_id, str) or not self.profile_id.strip(): + raise ValueError("motion limits profile_id must be a nonempty string") + for name in ("velocity_limits", "acceleration_limits"): + value = getattr(self, name) + if ( + not isinstance(value, torch.Tensor) + or value.ndim != 1 + or not value.numel() + or not value.is_floating_point() + or not bool(torch.isfinite(value).all()) + or not bool((value > 0).all()) + ): + raise ValueError(f"{name} must be a positive finite floating vector") + setattr( + self, name, value.detach().to(device="cpu", dtype=torch.float64).clone() + ) + if self.velocity_limits.shape != self.acceleration_limits.shape: + raise ValueError("motion limit vectors must use the same joint order") + + +def _plain(value: object) -> object: + if isinstance(value, Mapping): + return {key: _plain(item) for key, item in value.items()} + if is_dataclass(value): + return { + field.name: _plain(getattr(value, field.name)) for field in fields(value) + } + if isinstance(value, (tuple, list)): + return [_plain(item) for item in value] + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + return value + + +class GenerationRunner: + """Collect qualified qpos variants with one synchronous owner. + + The caller supplies explicit templates, a qualified initial-state profile, + motion collision model, actual-observation encoder, and task validator. + This runner supports handwritten free-motion sources and offline atomic + PickUp exports in ``env_rows`` / ``full_batch`` / synchronous LeRobot mode. + PickUp requires one shared full-state/contact validator for planning and + pure-sim execution. AtomicActionRuntime and contact-aware Gym execution + remain separate integrations. + + Planning failures do not reset physics. Ready candidates reserve episode + memory and target quota before restoration/execution. Only actual accepted + rollouts reach the sink, and only confirmed receipts increase coverage. + The runner is single-use and closes the supplied host and sink on exit. + ``generation_report.json`` preserves resolved settings, audit, and outcome. + + Args: + cfg: Strict job settings matching the supplied integration IDs. + host: Exclusive owner of the full simulator/Gym batch. + planner: Environment-row motion validation adapter for this robot. + executor: Concrete qpos rollout adapter for this host. + sink: Synchronous, local LeRobot writer with readback confirmation. + motion_limits: Trusted full-joint speed/acceleration policy. + max_write_retries: Extra synchronous submissions after a failed receipt. + clock: Monotonic clock used for the job wall-time budget. + """ + + def __init__( + self, + cfg: TrajectoryGenerationJobCfg, + host: FixedSceneHost, + planner: EnvRowMotionPlanner | PickUpMotionValidator, + executor: QposRolloutExecutor, + sink: LeRobotEpisodeSink, + *, + motion_limits: MotionLimitsProfile, + max_write_retries: int = 1, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self.cfg = TrajectoryGenerationJobCfg.from_mapping(cfg.to_dict()) + self.host, self.planner, self.executor, self.sink = ( + host, + planner, + executor, + sink, + ) + self.motion_limits = motion_limits.copy() + self.clock = clock + self._ran = False + if type(max_write_retries) is not int or max_write_retries < 0: + raise ValueError("max_write_retries must be a nonnegative integer") + self.max_write_retries = max_write_retries + contact_planner = isinstance(planner, PickUpMotionValidator) + if cfg.source.kind != "handwritten" and not contact_planner: + raise ValueError("this runner requires a handwritten qpos source") + if contact_planner and executor.contact_validator is not planner: + raise ValueError( + "PickUp planning and substep validation must share one validator" + ) + if ( + not contact_planner + and getattr(executor, "contact_validator", None) is not None + ): + raise ValueError( + "Contact execution requires complete PickUp planning validation" + ) + if contact_planner and cfg.augmentation.factors.timing.enabled: + raise ValueError("PickUp contact timing remains fixed") + if planner.robot is not host.adapter.robot or executor.host is not host: + raise ValueError("planner and executor must use the same owned host robot") + expected = { + "prepare profile": (cfg.reset.prepare_profile_id, host.profile.profile_id), + "initial-state tolerances": ( + cfg.reset.initial_state_tolerances_profile_id, + host.adapter.tolerances_profile_id, + ), + "motion limits": ( + cfg.validation.motion_limits_profile_id, + motion_limits.profile_id, + ), + "validator": (cfg.validation.validator_id, executor.validator_id), + "validation profile": ( + cfg.validation.profile_id, + executor.validation_profile_id, + ), + "sink": (cfg.persistence.sink, "lerobot"), + } + for label, (requested, actual) in expected.items(): + if requested != actual: + raise ValueError( + f"{label} ID {requested!r} does not match supplied {actual!r}" + ) + if len(motion_limits.velocity_limits) != len(host.adapter.robot.joint_names): + raise ValueError("motion limits must cover every robot joint") + if not math.isclose( + executor.control_dt * sink.fps, 1, rel_tol=1e-7, abs_tol=1e-9 + ): + raise ValueError("sink fps must match the authoritative control period") + cfg.validate_capabilities( + source_ids=(cfg.source.source_id,), + validator_ids=(executor.validator_id,), + profile_ids=( + host.profile.profile_id, + host.adapter.tolerances_profile_id, + motion_limits.profile_id, + executor.validation_profile_id, + ), + sink_ids=("lerobot",), + operators=("joint_residual", "retime"), + ) + self._episode_budget = executor.max_episode_bytes + if self._episode_budget > min( + sink.max_episode_bytes, cfg.persistence.pending_max_bytes + ): + raise ValueError( + "executor max_episode_bytes must fit both sink and pending byte limits" + ) + + def _validate_templates(self, cases, templates) -> None: + if len(cases) != self.host.adapter.sim.num_envs or len(templates) != len(cases): + raise ValueError("provide one case and template per physical row") + keys = [(case.scene_case_id, case.initial_state_id) for case in cases] + if len(set(keys)) != len(keys): + raise ValueError("per_env_case requires distinct case/initial-state pairs") + for template in templates: + if not isinstance(template, TrajectoryTemplate): + raise TypeError("templates must be explicit TrajectoryTemplate values") + if ( + template.source_id != self.cfg.source.source_id + or template.template_id != self.cfg.source.template_id + or template.validator_id != self.cfg.validation.validator_id + ): + raise ValueError( + "template source/template/validator IDs do not match the job" + ) + if template.joint_names != tuple(self.host.adapter.robot.joint_names): + raise ValueError( + "template joint names must match the complete robot order" + ) + if len(template.positions) < 2: + raise ValueError("an executable template needs at least two samples") + factors = self.cfg.augmentation.factors + needed = [] + if factors.spatial.enabled: + needed.append(factors.spatial.method) + if factors.timing.enabled: + needed.append("retime") + if set(needed) - set(template.allowed_operators): + raise ValueError( + "job enables operators not allowed by the source template" + ) + if not factors.timing.enabled and not torch.allclose( + template.dt[1:].double(), + torch.full_like(template.dt[1:].double(), self.executor.control_dt), + atol=1e-8, + rtol=0, + ): + raise ValueError( + "template timing must match the host clock or permit retime" + ) + + def _variant(self, template, limits, generator): + factors = self.cfg.augmentation.factors + value = template + metadata = {} + if factors.spatial.enabled: + value = joint_residual( + value, + joint_limits=limits, + normalized_scale=factors.spatial.joint_offset_scale, + generator=generator, + ) + metadata["joint_offset_scale"] = factors.spatial.joint_offset_scale + if factors.timing.enabled: + scales = factors.timing.duration_scales + scale = scales[int(torch.randint(len(scales), (), generator=generator))] + # Account for positions and arrival intervals before allocating retime output. + max_samples = self.cfg.execution.ready_max_bytes // ( + value.positions.element_size() * value.positions.shape[1] + + value.dt.element_size() + ) + value = retime( + value, + duration_scale=scale, + control_dt=self.executor.control_dt, + max_samples=max_samples, + ) + metadata["duration_scale"] = scale + return value, metadata + + def _quality(self, positions, times, reference, limits) -> ValidationResult: + ranges = (limits[:, 1] - limits[:, 0]).to( + device=positions.device, dtype=torch.float64 + ) + actual_length = float( + torch.linalg.vector_norm( + positions.double().diff(dim=0) / ranges, dim=1 + ).sum() + ) + reference_length = float( + torch.linalg.vector_norm( + reference.positions.to(positions.device).double().diff(dim=0) / ranges, + dim=1, + ).sum() + ) + duration = float(times[-1] - times[0]) + reference_duration = float(reference.dt.double().sum()) + length_ok = ( + actual_length + <= reference_length * self.cfg.validation.path_length_ratio_max + 1e-9 + ) + duration_ok = ( + duration + <= reference_duration * self.cfg.validation.duration_ratio_max + 1e-9 + ) + metrics = { + "path_length": actual_length, + "reference_path_length": reference_length, + "duration_s": duration, + "reference_duration_s": reference_duration, + } + return ValidationResult( + ( + ValidationCheck( + "motion_quality", + "passed" if length_ok and duration_ok else "failed", + "actual motion compared with the same-case reference", + metrics, + ), + ) + ) + + def _submit(self, session, episode) -> None: + for submission_id in range(self.max_write_retries + 1): + if submission_id: + episode, submission_id = session.retry_write(episode.commit_id) + try: + receipt = self.sink.submit(episode, submission_id=submission_id) + except Exception as error: + # This concrete synchronous sink raises input errors before any + # writes; all write/readback failures instead return a receipt. + session.apply_receipt( + CommitReceipt( + episode.episode_id, + episode.identity.candidate_id, + episode.identity.attempt_id, + "unsubmitted", + episode.commit_id, + episode.identity.scene_case_id, + confirmed=False, + error=f"{type(error).__name__}: {error}", + submission_id=submission_id, + ) + ) + raise + session.apply_receipt(receipt) + if receipt.confirmed: + return + raise RuntimeError(f"Episode persistence failed: {receipt.error}") + + def _write_report(self, report) -> None: + path = self.sink.root / "generation_report.json" + payload = json.dumps( + _plain(report), + ensure_ascii=False, + sort_keys=True, + indent=2, + allow_nan=False, + ).encode() + temporary = None + try: + with tempfile.NamedTemporaryFile( + dir=path.parent, prefix=".generation-report.", delete=False + ) as stream: + temporary = stream.name + stream.write(payload) + os.replace(temporary, path) + finally: + if temporary is not None: + Path(temporary).unlink(missing_ok=True) + + def run( + self, + cases: Sequence[SceneCase], + templates: Sequence[TrajectoryTemplate], + *, + should_stop: Callable[[], bool] | None = None, + ) -> Mapping[str, object]: + """Run one bounded collection job and freeze its final audit. + + Args: + cases: One caller-provided fixed case per simulator row. + templates: Explicit full-joint references in matching physical row order. + should_stop: Optional cancellation predicate checked between commands. + + Returns: + Counters, coverage, audit, resolved configuration, and ``target_reached``. + Exhausting a budget returns ``target_reached=False`` and a stop reason. + + Raises: + ValueError: If source or integration contracts are incompatible. + RuntimeError: If required validation, restoration, or persistence fails. + """ + if self._ran: + raise RuntimeError("GenerationRunner is single-use") + self._ran = True + session = GenerationSession(self.cfg, clock=self.clock) + error = None + cancelled = False + cases, templates = tuple(cases), tuple(templates) + try: + self._validate_templates(cases, templates) + binding = self.host.acquire_case(cases) + snapshots = self.host.snapshots(binding) + if isinstance(self.planner, PickUpMotionValidator): + world_ids = updated_ids = set(self.planner.collision_world_entity_ids) + else: + world_ids = set( + self.planner.motion_generator.collision_world_entity_ids + ) + updated_ids = set( + self.planner.motion_generator.dynamic_collision_entity_ids + ) + if any( + set(snapshot.entity_poses) - (world_ids & updated_ids) + for snapshot in snapshots + ): + raise ValueError( + "generation collision world must update every rigid object under its physical UID" + ) + limits = self.host.adapter.robot.get_qpos_limits().detach().clone() + for case, template, row_limits in zip(cases, templates, limits): + session.register_case( + case, row_limits, joint_names=template.joint_names + ) + plan_checks = {} + reference_by_candidate = {} + ready_by_row = set() + row_offset = 0 + needs_restore = False + + def stop() -> bool: + return bool( + session.stop_reason or (should_stop is not None and should_stop()) + ) + + while not stop(): + order = [ + (row_offset + offset) % len(cases) for offset in range(len(cases)) + ] + for row in order: + if ( + row in ready_by_row + or stop() + or session.snapshot()["counts"]["ready"] + >= self.cfg.execution.ready_high_watermark + ): + continue + if ( + session.snapshot()["counts"]["proposed"] + >= self.cfg.collection.max_proposals + ): + break + case, reference = cases[row], templates[row] + family = None + if not self.cfg.augmentation.factors.spatial.enabled: + family = ( + "geometry_" + + hashlib.sha256( + json.dumps( + ( + case.scene_case_id, + case.initial_state_id, + reference.source_id, + reference.source_revision, + reference.template_id, + ) + ).encode() + ).hexdigest() + ) + identity, generator = session.propose( + case.scene_case_id, + case.initial_state_id, + source_id=reference.source_id, + source_revision=reference.source_revision, + template_id=reference.template_id, + operator_id="configured_factors", + geometry_family_id=family, + ) + try: + variant, factors = self._variant( + reference, limits[row], generator + ) + batch = CandidateTrajectoryBatch( + variant.positions.unsqueeze(0), + variant.dt.unsqueeze(0), + torch.tensor([len(variant.positions)], dtype=torch.int64), + (identity,), + variant.joint_names, + (variant.phases,), + (factors,), + torch.tensor([row], dtype=torch.int64), + ) + collision = self.planner.validate_qpos(batch, snapshots)[0] + dynamics = validate_motion_limits( + variant, + velocity_limits=self.motion_limits.velocity_limits, + acceleration_limits=self.motion_limits.acceleration_limits, + ) + validation = ValidationResult( + (*collision.checks, *dynamics.checks) + ) + try: + session.add_planned(batch, validation) + except ValueError: + if any( + check.status in {"unavailable", "not_run"} + for check in validation.checks + ): + raise RuntimeError( + f"Required planning capability is unavailable: {validation.checks}" + ) from None + raise + except (ValueError, BufferError) as failure: + session.release( + identity, reason=f"planning rejected: {failure}" + ) + continue + plan_checks[identity.candidate_id] = validation + reference_by_candidate[identity.candidate_id] = reference + ready_by_row.add(row) + if stop(): + break + selected = [None] * len(cases) + for row in order: + case = cases[row] + selected[row] = session.take_ready( + case.scene_case_id, + case.initial_state_id, + episode_byte_budget=self._episode_budget, + ) + if selected[row] is not None: + ready_by_row.discard(row) + assigned = [ + candidate for candidate in selected if candidate is not None + ] + if not assigned: + if session.stop_reason: + break + if ( + session.snapshot()["counts"]["proposed"] + >= self.cfg.collection.max_proposals + ): + break + continue + if stop(): + break + if needs_restore: + binding = self.host.restore_initial() + verified = self.host.verify_initial(binding) + if not verified.accepted: + raise RuntimeError(f"Initial state mismatch: {verified.checks}") + if stop(): + break + ids = { + candidate.identities[0].candidate_id: session.episode_ids( + candidate.identities[0] + ) + for candidate in assigned + } + episodes = self.executor.execute( + binding, + selected, + ids, + on_started=session.mark_rollout_started, + should_stop=stop, + ) + needs_restore = True + if len(episodes) != len(cases): + raise ValueError( + "executor results must preserve physical row order" + ) + for row, (candidate, episode) in enumerate(zip(selected, episodes)): + if candidate is None: + if episode is not None: + raise ValueError( + "executor produced evidence for an inactive row" + ) + continue + identity = candidate.identities[0] + if episode is None: + plan_checks.pop(identity.candidate_id, None) + reference_by_candidate.pop(identity.candidate_id, None) + session.release( + identity, + reason=self.executor.last_failures.get( + identity.candidate_id, + "execution produced no complete evidence", + ), + ) + continue + if episode.identity != identity: + raise ValueError( + "executor evidence belongs to a different candidate" + ) + reference = reference_by_candidate.pop(identity.candidate_id) + qpos = episode.observations["joint_positions"] + quality = self._quality( + qpos, episode.timestamps, reference, limits[row] + ) + measured = replace( + reference, + positions=qpos, + dt=torch.cat( + (episode.timestamps.new_zeros(1), episode.timestamps.diff()) + ), + phases=episode.phases, + ) + dynamic = validate_motion_limits( + measured, + velocity_limits=self.motion_limits.velocity_limits, + acceleration_limits=self.motion_limits.acceleration_limits, + ) + # Revalidate measured joints. Free motion uses the captured + # fixed world; PickUp uses actual target poses, with root + # and all other objects verified by the executor. + actual_snapshots = list(snapshots) + actual_snapshots[row] = replace( + snapshots[row], joint_positions=qpos[0] + ) + actual_batch = CandidateTrajectoryBatch( + qpos.unsqueeze(0), + measured.dt.unsqueeze(0), + torch.tensor([len(qpos)], dtype=torch.int64), + (identity,), + candidate.joint_names, + (episode.phases,), + source_row_indices=torch.tensor([row], dtype=torch.int64), + ) + if isinstance(self.planner, PickUpMotionValidator): + actual_collision = self.planner.validate_episode( + episode, + actual_snapshots[row], + row=row, + ) + else: + actual_collision = self.planner.validate_qpos( + actual_batch, + actual_snapshots, + )[0] + all_checks = ( + *tuple( + replace(check, check_id="planned_" + check.check_id) + for check in plan_checks.pop(identity.candidate_id).checks + ), + *actual_collision.checks, + *episode.validation.checks, + *quality.checks, + *tuple( + ValidationCheck( + "actual_" + check.check_id, + check.status, + check.detail, + check.metrics, + ) + for check in dynamic.checks + ), + ) + accepted = replace(episode, validation=ValidationResult(all_checks)) + if session.accept_episode(accepted): + self._submit(session, accepted) + row_offset = (order[0] + 1) % len(cases) + cancelled = ( + should_stop is not None + and should_stop() + and session.stop_reason is None + ) + except BaseException as failure: + error = failure + finally: + try: + for receipt in self.sink.drain(): + session.apply_receipt(receipt) + except BaseException as failure: + if error is None: + error = failure + for identity, state, _ in session.snapshot()["audit"]: + if state in { + "proposed", + "ready", + "assigned", + "running", + "write_failed", + }: + session.release( + identity, + reason="job stopped" if error is None else "job failed", + ) + for close in (self.sink.close, self.host.close): + try: + close() + except BaseException as failure: + if error is None: + error = failure + report = { + **dict(session.snapshot()), + "configuration": self.cfg.to_dict(), + "motion_limits": self.motion_limits, + "cases": cases, + "version": __version__, + "target_reached": session.snapshot()["counts"]["committed"] + >= self.cfg.collection.target_committed_episodes, + "cancelled": cancelled, + "error": None if error is None else f"{type(error).__name__}: {error}", + } + try: + self._write_report(report) + except BaseException as failure: + if error is None: + error = failure + else: + add_note = getattr(error, "add_note", None) + if callable(add_note): + add_note( + f"Writing generation_report.json also failed: {failure}" + ) + else: + error.__context__ = failure + if error is not None: + raise error + return report diff --git a/embodichain/lab/trajectory_generation/sinks.py b/embodichain/lab/trajectory_generation/sinks.py new file mode 100644 index 000000000..e279f3781 --- /dev/null +++ b/embodichain/lab/trajectory_generation/sinks.py @@ -0,0 +1,553 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Synchronous, host-independent persistence of accepted expert episodes.""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import re +import shutil +import tempfile +from collections.abc import Mapping +from dataclasses import asdict +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CommitReceipt, + ExpertEpisode, +) + +__all__ = ["LeRobotEpisodeSink"] + +_NUMERIC_DTYPES = {torch.float32, torch.float64, torch.int32, torch.int64, torch.uint8} +_FEATURE_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*") + + +def _plain(value: Any) -> Any: + """Convert immutable contract metadata to JSON-native containers.""" + if isinstance(value, Mapping): + return {key: _plain(item) for key, item in value.items()} + if isinstance(value, (set, frozenset)): + return [_plain(item) for item in sorted(value)] + if isinstance(value, (tuple, list)): + return [_plain(item) for item in value] + return value + + +def _json_bytes(value: Any) -> bytes: + return json.dumps( + _plain(value), sort_keys=True, separators=(",", ":"), allow_nan=False + ).encode("utf-8") + + +def _atomic_bytes(path: Path, data: bytes) -> None: + """Replace one required metadata file only after its complete write.""" + temporary: str | None = None + try: + with tempfile.NamedTemporaryFile( + dir=path.parent, prefix=f".{path.name}.", delete=False + ) as stream: + temporary = stream.name + stream.write(data) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + finally: + if temporary is not None: + Path(temporary).unlink(missing_ok=True) + + +class LeRobotEpisodeSink: + """Seal and read back one LeRobot dataset shard per accepted episode. + + Each shard contains exactly ``T`` causal training frames. Required sidecars + retain the terminal observation, all ``T+1`` measured timestamps, candidate + lineage, validation, phases, and user metadata. Only shards listed in + ``manifest.json`` are committed. Every successful receipt follows writer + finalization, dataset/media decoding, sidecar verification, and manifest + replacement/readback. This is a synchronous, single-caller sink; ``drain`` + has no deferred receipts. + + Numeric observations must be nonempty vectors or matrices of float32, + float64, int32, int64, or uint8. Matrices are flattened in C order into + LeRobot vectors; ``observation_shapes`` in episode.json retains their source + layout, and terminal.npz preserves the original shape. RGB observations + must be uint8 arrays shaped ``(T+1,H,W,3)``. + Feature names are preserved below ``observation.``; already prefixed names + are retained. Unsupported fields are rejected before any episode write. + Actions must be float32 or float64 vectors. Training timestamps are the + nominal relative LeRobot clock; exact measured timestamps, including their + origin, are authoritative in the required sidecar. + + One episode is the bounded shard unit. Retrying an identical commit reuses + a readable sealed shard; incomplete uncommitted data is rebuilt at the same + path. Changed payloads under the same commit ID are rejected. The sink owns + a new/empty output directory and does not implement process restart or + power-loss recovery. Do not concurrently modify its files. + + Args: + root: New or empty output directory for the collection manifest/shards. + fps: Integer control frequency. All measured intervals and relative + timestamps must agree with this clock within ``timestamp_tolerance``. + repo_id: Local LeRobot repository identity; nothing is uploaded. + max_episode_bytes: Maximum raw tensor plus metadata bytes per episode. + timestamp_tolerance: Absolute seconds allowed between the measured and + nominal clocks; relative tolerance is zero. + """ + + def __init__( + self, + root: str | Path, + *, + fps: int, + repo_id: str = "embodichain/trajectory-generation", + max_episode_bytes: int = 256 * 1024 * 1024, + timestamp_tolerance: float = 1e-6, + ) -> None: + if type(fps) is not int or fps < 1: + raise ValueError("fps must be a positive integer.") + if type(max_episode_bytes) is not int or max_episode_bytes < 1: + raise ValueError("max_episode_bytes must be a positive integer.") + if not isinstance(repo_id, str) or not repo_id: + raise ValueError("repo_id must be a nonempty string.") + if ( + isinstance(timestamp_tolerance, bool) + or not isinstance(timestamp_tolerance, (int, float)) + or not math.isfinite(timestamp_tolerance) + or timestamp_tolerance < 0 + ): + raise ValueError("timestamp_tolerance must be a finite nonnegative number.") + self.root = Path(root).resolve() + self.fps = fps + self.repo_id = repo_id + self.max_episode_bytes = max_episode_bytes + self.timestamp_tolerance = float(timestamp_tolerance) + self._closed = False + self._fingerprints: dict[str, str] = {} + self._committed: dict[str, dict[str, Any]] = {} + self.root.mkdir(parents=True, exist_ok=True) + if any(self.root.iterdir()): + raise ValueError( + "LeRobotEpisodeSink requires a new or empty root directory." + ) + self._lock = self.root / ".writer.lock" + self._lock.touch(exist_ok=False) + + def _prepare(self, episode: ExpertEpisode) -> tuple[ExpertEpisode, dict, dict, str]: + """Validate, bound, and own the complete submitted evidence.""" + if not isinstance(episode, ExpertEpisode): + raise TypeError("episode must be an ExpertEpisode.") + if not episode.validation.accepted: + raise ValueError("Only accepted episodes may enter the expert dataset.") + metadata = { + "identity": asdict(episode.identity), + "episode_id": episode.episode_id, + "commit_id": episode.commit_id, + "action_representation": episode.action_representation, + "validation": [ + { + "check_id": check.check_id, + "status": check.status, + "detail": check.detail, + "metrics": dict(check.metrics), + } + for check in episode.validation.checks + ], + "phases": [asdict(phase) for phase in episode.phases], + "metadata": _plain(episode.metadata), + "fps": self.fps, + "steps": episode.actions.shape[0], + "observation_features": {}, + "observation_shapes": {}, + } + # Reject unknown layouts before the CPU ownership copy or filesystem work. + features = {} + for key, value in episode.observations.items(): + if not _FEATURE_NAME.fullmatch(key): + raise ValueError(f"Unsupported observation feature name: {key!r}.") + name = key if key.startswith("observation.") else f"observation.{key}" + if name in features: + raise ValueError(f"Observation feature names collide at {name!r}.") + if ( + value.ndim in (2, 3) + and min(value.shape[1:]) > 0 + and value.dtype in _NUMERIC_DTYPES + ): + features[name] = { + "dtype": str(value.dtype).removeprefix("torch."), + "shape": (math.prod(value.shape[1:]),), + "names": None, + } + elif ( + value.ndim == 4 + and value.dtype == torch.uint8 + and value.shape[-1] == 3 + and min(value.shape[1:3]) > 0 + ): + features[name] = { + "dtype": "image", + "shape": (3, *value.shape[1:3]), + "names": ["channel", "height", "width"], + } + else: + raise ValueError( + f"Unsupported observation shape/dtype for {key!r}: {tuple(value.shape)}, {value.dtype}." + ) + metadata["observation_features"][key] = name + metadata["observation_shapes"][key] = list(value.shape[1:]) + if episode.actions.dtype not in {torch.float32, torch.float64}: + raise ValueError("Actions must use float32 or float64.") + features["action"] = { + "dtype": str(episode.actions.dtype).removeprefix("torch."), + "shape": tuple(episode.actions.shape[1:]), + "names": None, + } + metadata["features"] = features + encoded = _json_bytes(metadata) + tensors = (episode.actions, episode.timestamps, *episode.observations.values()) + byte_count = len(encoded) + sum( + value.numel() * value.element_size() for value in tensors + ) + if byte_count > self.max_episode_bytes: + raise ValueError("Episode exceeds max_episode_bytes.") + owned = ExpertEpisode( + identity=episode.identity, + observations={ + key: value.detach().cpu() for key, value in episode.observations.items() + }, + actions=episode.actions.detach().cpu(), + timestamps=episode.timestamps.detach().cpu(), + action_representation=episode.action_representation, + validation=episode.validation, + episode_id=episode.episode_id, + commit_id=episode.commit_id, + phases=episode.phases, + metadata=episode.metadata, + ) + times = owned.timestamps.to(torch.float64) + relative = times - times[0] + nominal = torch.arange(times.numel(), dtype=torch.float64) / self.fps + if not torch.allclose( + relative, nominal, atol=self.timestamp_tolerance, rtol=0 + ) or not torch.allclose( + times.diff(), + torch.full_like(times[1:], 1 / self.fps), + atol=self.timestamp_tolerance, + rtol=0, + ): + raise ValueError( + "Measured timestamps do not match the configured fixed control clock." + ) + digest = hashlib.sha256(encoded) + for name, value in [ + ("actions", owned.actions), + ("timestamps", owned.timestamps), + *sorted(owned.observations.items()), + ]: + digest.update(_json_bytes([name, str(value.dtype), list(value.shape)])) + digest.update(value.contiguous().numpy().tobytes()) + return owned, features, json.loads(encoded), digest.hexdigest() + + def _write_dataset( + self, path: Path, episode: ExpertEpisode, features: dict, metadata: dict + ) -> None: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + dataset = LeRobotDataset.create( + repo_id=self.repo_id, + root=path, + fps=self.fps, + features=features, + use_videos=False, + metadata_buffer_size=1, + image_writer_threads=0, + image_writer_processes=0, + ) + try: + for index, action in enumerate(episode.actions): + frame = { + "action": action.numpy(), + "task": str( + episode.metadata.get("task", episode.identity.source_id) + ), + } + for key, value in episode.observations.items(): + name = metadata["observation_features"][key] + sample = value[index] + if features[name]["dtype"] == "image": + sample = sample.permute(2, 0, 1) + else: + sample = sample.reshape(-1) + frame[name] = sample.numpy() + dataset.add_frame(frame) + # LeRobot 0.4.4 validates (1,) numeric arrays but serializes them as + # scalar HF Values. Match that schema without NumPy 2 scalar casts. + for name, feature in features.items(): + if feature["shape"] == (1,) and feature["dtype"] != "image": + dataset.episode_buffer[name] = [ + value.reshape(-1)[0] for value in dataset.episode_buffer[name] + ] + dataset.save_episode() + finally: + dataset.finalize() + + def _verify_dataset( + self, path: Path, episode: ExpertEpisode, features: dict, metadata: dict + ) -> None: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + import pyarrow.parquet as pq + + # Check local completeness before LeRobot's constructor can attempt a + # Hub fallback for missing metadata. Footer reads also require sealing. + required = [ + path / "meta/info.json", + path / "meta/stats.json", + path / "meta/tasks.parquet", + ] + data_files = list((path / "data").rglob("*.parquet")) + episode_files = list((path / "meta/episodes").rglob("*.parquet")) + if ( + not all(file.is_file() for file in required) + or not data_files + or not episode_files + ): + raise ValueError("LeRobot shard is not sealed and locally complete.") + for file in [*data_files, *episode_files, required[-1]]: + pq.ParquetFile(file).metadata + dataset = LeRobotDataset(repo_id=self.repo_id, root=path, download_videos=False) + if ( + dataset.num_episodes != 1 + or len(dataset) != len(episode.actions) + or dataset.fps != self.fps + ): + raise ValueError( + "LeRobot shard has incorrect episode, frame, or clock metadata." + ) + # LeRobot's torch row transform constructs Python float lists using the + # default float32 dtype. Read declared float64 columns from their sealed + # Arrow representation to verify precision without that reader cast. + precise = { + name: [] + for name, feature in features.items() + if feature["dtype"] == "float64" + } + if precise: + for file in sorted(data_files): + table = pq.read_table(file, columns=list(precise)) + for name in precise: + precise[name].extend(table[name].to_pylist()) + for index, action in enumerate(episode.actions): + frame = dataset[index] # Decodes every required image as well. + observed_action = ( + ( + torch.as_tensor(precise["action"][index], dtype=action.dtype) + if "action" in precise + else frame["action"] + ) + .reshape(action.shape) + .to(action.dtype) + ) + if not torch.equal(observed_action, action): + raise ValueError( + "LeRobot action readback differs from submitted commands." + ) + if not torch.equal( + frame["timestamp"], + torch.tensor(index / self.fps, dtype=frame["timestamp"].dtype), + ): + raise ValueError( + "LeRobot frame timestamps do not match the fixed clock." + ) + for key, value in episode.observations.items(): + name = metadata["observation_features"][key] + current = frame[name] + expected = value[index] + if features[name]["dtype"] == "image": + current = (current * 255).round().to(torch.uint8).permute(1, 2, 0) + else: + if name in precise: + current = torch.as_tensor( + precise[name][index], dtype=expected.dtype + ) + current = current.reshape(expected.shape).to(expected.dtype) + if not torch.equal(current, expected): + raise ValueError( + f"LeRobot observation readback differs for {key!r}." + ) + + def _write_evidence( + self, path: Path, episode: ExpertEpisode, metadata: dict + ) -> None: + _atomic_bytes(path / "episode.json", _json_bytes(metadata)) + temporary: str | None = None + try: + with tempfile.NamedTemporaryFile( + dir=path, prefix=".terminal.", delete=False + ) as stream: + temporary = stream.name + arrays = {"timestamps": episode.timestamps.numpy()} + arrays.update( + { + f"observation.{key}": value[-1].numpy() + for key, value in episode.observations.items() + } + ) + np.savez(stream, **arrays) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path / "terminal.npz") + finally: + if temporary is not None: + Path(temporary).unlink(missing_ok=True) + + def _verify_evidence( + self, path: Path, episode: ExpertEpisode, metadata: dict + ) -> None: + if json.loads((path / "episode.json").read_text()) != metadata: + raise ValueError("Episode lineage/validation metadata readback differs.") + expected = {"timestamps": episode.timestamps.numpy()} + expected.update( + { + f"observation.{key}": value[-1].numpy() + for key, value in episode.observations.items() + } + ) + with np.load(path / "terminal.npz", allow_pickle=False) as arrays: + if set(arrays.files) != set(expected): + raise ValueError("Terminal evidence is incomplete.") + for key, value in expected.items(): + if ( + arrays[key].dtype != value.dtype + or arrays[key].shape != value.shape + or not np.array_equal(arrays[key], value) + ): + raise ValueError(f"Terminal evidence readback differs for {key!r}.") + + def _write_manifest(self, record: dict) -> None: + updated = {**self._committed, record["commit_id"]: record} + manifest = { + "format": "embodichain.expert_shards", + "fps": self.fps, + "episodes": list(updated.values()), + } + _atomic_bytes(self.root / "manifest.json", _json_bytes(manifest)) + self._verify_manifest(updated) + self._committed = updated + + def _verify_manifest(self, records: dict) -> None: + expected = { + "format": "embodichain.expert_shards", + "fps": self.fps, + "episodes": list(records.values()), + } + if json.loads((self.root / "manifest.json").read_text()) != expected: + raise ValueError("Collection manifest readback differs.") + + def submit( + self, episode: ExpertEpisode, *, submission_id: int = 0 + ) -> CommitReceipt: + """Persist accepted evidence and return its final synchronous receipt. + + Invalid/unsupported input and changed payloads under an existing commit + raise before episode writes. Persistence failures return an unconfirmed + receipt. Retrying uses the same ``commit_id`` and a new submission ID. + """ + if self._closed: + raise RuntimeError("LeRobotEpisodeSink is closed.") + if type(submission_id) is not int or submission_id < 0: + raise ValueError("submission_id must be a nonnegative integer.") + owned, features, metadata, fingerprint = self._prepare(episode) + previous = self._fingerprints.setdefault(owned.commit_id, fingerprint) + if previous != fingerprint: + raise ValueError( + "The same commit_id cannot identify a changed episode payload." + ) + path = self.root / hashlib.sha256(owned.commit_id.encode()).hexdigest() + dataset_path = path / "dataset" + confirmed, error = False, "" + try: + if owned.commit_id in self._committed: + self._verify_dataset(dataset_path, owned, features, metadata) + self._verify_evidence(path, owned, metadata) + self._verify_manifest(self._committed) + return CommitReceipt( + episode_id=owned.episode_id, + candidate_id=owned.identity.candidate_id, + attempt_id=owned.identity.attempt_id, + storage_id=str(path), + commit_id=owned.commit_id, + scene_case_id=owned.identity.scene_case_id, + submission_id=submission_id, + ) + path.mkdir(exist_ok=True) + if dataset_path.exists(): + try: + self._verify_dataset(dataset_path, owned, features, metadata) + except Exception: + if owned.commit_id in self._committed: + raise + shutil.rmtree(dataset_path) + if not dataset_path.exists(): + self._write_dataset(dataset_path, owned, features, metadata) + self._verify_dataset(dataset_path, owned, features, metadata) + self._write_evidence(path, owned, metadata) + self._verify_evidence(path, owned, metadata) + self._write_manifest( + { + "commit_id": owned.commit_id, + "episode_id": owned.episode_id, + "candidate_id": owned.identity.candidate_id, + "fingerprint": fingerprint, + "shard": path.name, + } + ) + confirmed = True + except Exception as exception: + error = f"{type(exception).__name__}: {exception}"[:1024] + return CommitReceipt( + episode_id=owned.episode_id, + candidate_id=owned.identity.candidate_id, + attempt_id=owned.identity.attempt_id, + storage_id=str(path), + commit_id=owned.commit_id, + scene_case_id=owned.identity.scene_case_id, + confirmed=confirmed, + error=error, + submission_id=submission_id, + ) + + def drain(self) -> tuple[CommitReceipt, ...]: + """Return no deferred work: each submit already returns its final receipt.""" + return () + + def close(self) -> None: + """Release the writer lease; all successful submissions are already sealed.""" + self._closed = True + self._lock.unlink(missing_ok=True) + + def __enter__(self) -> LeRobotEpisodeSink: + return self + + def __exit__(self, *args: object) -> None: + self.close() diff --git a/embodichain_tasks/embodichain_tasks/manipulation/tableware/blocks_ranking_rgb.py b/embodichain_tasks/embodichain_tasks/manipulation/tableware/blocks_ranking_rgb.py index 1eaa2d3ef..59881f094 100644 --- a/embodichain_tasks/embodichain_tasks/manipulation/tableware/blocks_ranking_rgb.py +++ b/embodichain_tasks/embodichain_tasks/manipulation/tableware/blocks_ranking_rgb.py @@ -94,7 +94,7 @@ def _initialize_atomic_actions(self) -> None: ControlPartCommandProfile, ObjectSemantics, ) - from embodichain.lab.sim.planners import ( + from embodichain.lab.sim.motion.planners import ( MotionGenCfg, MotionGenerator, ToppraPlannerCfg, diff --git a/embodichain_tasks/embodichain_tasks/manipulation/tableware/stack_blocks_two.py b/embodichain_tasks/embodichain_tasks/manipulation/tableware/stack_blocks_two.py index 47d5fd37e..e1382ee44 100644 --- a/embodichain_tasks/embodichain_tasks/manipulation/tableware/stack_blocks_two.py +++ b/embodichain_tasks/embodichain_tasks/manipulation/tableware/stack_blocks_two.py @@ -74,7 +74,7 @@ def _initialize_atomic_actions(self) -> None: ControlPartCommandProfile, ObjectSemantics, ) - from embodichain.lab.sim.planners import ( + from embodichain.lab.sim.motion.planners import ( MotionGenCfg, MotionGenerator, ToppraPlannerCfg, diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index e8271088e..2c7034364 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -24,7 +24,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.solvers import PinkSolverCfg, PytorchSolverCfg +from embodichain.lab.sim.motion.solvers import PinkSolverCfg, PytorchSolverCfg from embodichain.lab.sim.objects import ( GizmoCfg, create_robot_ik_gizmo_controller, diff --git a/examples/sim/planners/curobo_planner.py b/examples/sim/motion/planners/curobo_planner.py similarity index 98% rename from examples/sim/planners/curobo_planner.py rename to examples/sim/motion/planners/curobo_planner.py index 908f8619a..b910f6e15 100644 --- a/examples/sim/planners/curobo_planner.py +++ b/examples/sim/motion/planners/curobo_planner.py @@ -26,9 +26,9 @@ Run from the repository root:: - python examples/sim/planners/curobo_planner.py --headless - python examples/sim/planners/curobo_planner.py --headless --num_envs 4 - python examples/sim/planners/curobo_planner.py --headless --device cuda:1 + python examples/sim/motion/planners/curobo_planner.py --headless + python examples/sim/motion/planners/curobo_planner.py --headless --num_envs 4 + python examples/sim/motion/planners/curobo_planner.py --headless --device cuda:1 Requirements: an NVIDIA CUDA device and the CUDA-matched cuRobo V2 source package installed in the active environment. Installation instructions: @@ -48,7 +48,7 @@ # package, so this example exercises the current code. The demo relies on the # cuRobo adapter's URDF-based robot-YAML auto-generation, which lives in the # source tree and may not be present in an older installed copy. -_REPO_ROOT = Path(__file__).resolve().parents[3] +_REPO_ROOT = Path(__file__).resolve().parents[4] if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) @@ -67,8 +67,8 @@ from embodichain.data import get_data_path from embodichain.lab.sim.cfg import RenderCfg, RigidBodyAttributesCfg from embodichain.lab.sim.objects import RigidObjectCfg, Robot, RigidObject -from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator -from embodichain.lab.sim.planners.curobo.curobo_planner import ( +from embodichain.lab.sim.motion.planners import MotionGenCfg, MotionGenerator +from embodichain.lab.sim.motion.planners.curobo.curobo_planner import ( CuroboPlanOptions, CuroboPlannerCfg, CuroboWorldCfg, diff --git a/examples/sim/planners/neural_planner.py b/examples/sim/motion/planners/neural_planner.py similarity index 97% rename from examples/sim/planners/neural_planner.py rename to examples/sim/motion/planners/neural_planner.py index 9ce5a18db..5f205a218 100644 --- a/examples/sim/planners/neural_planner.py +++ b/examples/sim/motion/planners/neural_planner.py @@ -18,8 +18,8 @@ From the repository root:: - python examples/sim/planners/neural_planner.py --headless - python examples/sim/planners/neural_planner.py --headless --device cuda:1 + python examples/sim/motion/planners/neural_planner.py --headless + python examples/sim/motion/planners/neural_planner.py --headless --device cuda:1 """ from __future__ import annotations @@ -37,7 +37,7 @@ from embodichain.lab.sim.cfg import MarkerCfg, RenderCfg from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.robots.franka_panda import FrankaPandaCfg -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( MotionGenCfg, MotionGenOptions, MotionGenerator, @@ -45,7 +45,7 @@ NeuralPlannerCfg, PlanState, ) -from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions +from embodichain.lab.sim.motion.planners.neural_planner import NeuralPlanOptions def parse_args() -> argparse.Namespace: diff --git a/examples/sim/solvers/differential_solver.py b/examples/sim/motion/solvers/differential_solver.py similarity index 100% rename from examples/sim/solvers/differential_solver.py rename to examples/sim/motion/solvers/differential_solver.py diff --git a/examples/sim/solvers/neural_ik_solver.py b/examples/sim/motion/solvers/neural_ik_solver.py similarity index 99% rename from examples/sim/solvers/neural_ik_solver.py rename to examples/sim/motion/solvers/neural_ik_solver.py index 5d2b9a83f..069108163 100644 --- a/examples/sim/solvers/neural_ik_solver.py +++ b/examples/sim/motion/solvers/neural_ik_solver.py @@ -29,7 +29,7 @@ from embodichain.lab.sim.cfg import MarkerCfg from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.robots.franka_panda import FrankaPandaCfg -from embodichain.lab.sim.solvers import NeuralIKSolverCfg +from embodichain.lab.sim.motion.solvers import NeuralIKSolverCfg from embodichain.lab.visualization import ( add_viser_args_to_parser, visualization_cfg_from_args, diff --git a/examples/sim/solvers/opw_solver.py b/examples/sim/motion/solvers/opw_solver.py similarity index 100% rename from examples/sim/solvers/opw_solver.py rename to examples/sim/motion/solvers/opw_solver.py diff --git a/examples/sim/solvers/pink_solver.py b/examples/sim/motion/solvers/pink_solver.py similarity index 100% rename from examples/sim/solvers/pink_solver.py rename to examples/sim/motion/solvers/pink_solver.py diff --git a/examples/sim/solvers/pinocchio_solver.py b/examples/sim/motion/solvers/pinocchio_solver.py similarity index 100% rename from examples/sim/solvers/pinocchio_solver.py rename to examples/sim/motion/solvers/pinocchio_solver.py diff --git a/examples/sim/solvers/pytorch_solver.py b/examples/sim/motion/solvers/pytorch_solver.py similarity index 100% rename from examples/sim/solvers/pytorch_solver.py rename to examples/sim/motion/solvers/pytorch_solver.py diff --git a/examples/sim/solvers/srs_solver.py b/examples/sim/motion/solvers/srs_solver.py similarity index 100% rename from examples/sim/solvers/srs_solver.py rename to examples/sim/motion/solvers/srs_solver.py diff --git a/examples/sim/motion/trajectory_generation/README.md b/examples/sim/motion/trajectory_generation/README.md new file mode 100644 index 000000000..b12e4df0f --- /dev/null +++ b/examples/sim/motion/trajectory_generation/README.md @@ -0,0 +1,97 @@ +# 固定场景轨迹生成与增强示例 + +## 并行抓取固定方块：姿态和路径增强 + +`cube_grasp_parallel.py` 在四个同步仿真环境中，让 UR5 + PGI 夹爪抓取同样的 5 cm 方块。机器人和方块的初始位姿相同，方块受重力和接触作用，夹起后能够移动。 + +在已安装 EmbodiChain 仿真依赖、UR5/PGI 资产及 TOPPRA 的环境中运行。视频使用离屏渲染，无需桌面窗口；此例使用 UR IK 和 `ik_interp`，不需要 cuRobo 后端。 + +```bash +python -m pip install 'imageio[ffmpeg]' +python examples/sim/motion/trajectory_generation/cube_grasp_parallel.py \ + --output /tmp/cube-grasp-parallel +``` + +输出目录必须不存在或为空。`--seed` 默认为 `13`，控制路径残差的局部随机种子；`--cuda-device` 默认为 `0`，选择渲染显卡。 + +| 四宫格位置 | 抓取朝向 | 自由运动路径 | +|---|---|---| +| 左上 | 0° 参考抓取 | 参考路径 | +| 右上 | 0° 参考抓取 | 增强路径 | +| 左下 | 绕方块局部 Z 轴旋转 90° | 为新抓取朝向重新规划的参考路径 | +| 右下 | 绕方块局部 Z 轴旋转 90° | 增强路径 | + +姿态增强调用 `rotate_grasp_about_object_axis`，利用方块的四分之一转对称性生成新的 TCP 目标；每个目标都重新求解 IK。路径增强调用 `joint_residual`，只改变到预抓取位置之前的自由运动段，保留起终点、夹爪命令及后续接近、闭合和抬升段。现有 `MoveEndEffector` 与 `PickUp` 技能负责动作规划。 + +四路在同一个 `SimulationManager`、同一个物理时钟下执行。每个控制时刻一起采集四个相机画面，再拼成四宫格；彩色线是实测 TCP 轨迹，标题显示抓取角度、运动阶段和方块抬升高度。 + +输出文件： + +- `preview.mp4`：1280×1056、20 fps 的同步对比视频，包含初始与终态画面。 +- `rollout.npz`：四路参考/指令/实测关节轨迹、实测 TCP 和方块位姿、抓取目标与统一时间戳。轨迹数组以 `(环境, 时间, ...)` 排列；参考轨迹不含末尾额外保持段。 +- `report.json`：增强参数、阶段边界、路径差异、各路抓取验收结果。末尾保持段要求抬升至少 12 cm、方块相对 TCP 的位置漂移不超过 1 cm，且始终靠近 TCP。 + +默认参数的实测结果为四路全部成功，保持期间抬升约 17.4–17.7 cm，两组增强/参考路径的最大 TCP 间距约 17 cm。物理运动为 13.55 秒，加上终态画面的播放间隔，272 帧视频时长为 13.60 秒。改变种子后仍会重新验收，未全部成功时返回退出码 1。 + +这是抓取增强的物理可视化示例，不写入专家 LeRobot 数据。需要碰撞、真实接触检查和多轮采集时，使用下方 `cube_pickup_collection.py`。 + +## 并行抓取的专家数据采集 + +`cube_pickup_collection.py` 复用上面的四路 cube 场景和原子动作规划，接入 `FixedSceneHost → GenerationRunner/GenerationSession → LeRobotEpisodeSink`。四路各自生成自由段残差，使用 0°/90° 抓取目标；每轮结束恢复整批初态，再继续生成，直到已确认提交数量达到目标或预算耗尽。 + +```bash +python -m pip install -e '.[trajectory-generation]' 'imageio[ffmpeg]' +python examples/sim/motion/trajectory_generation/cube_pickup_collection.py \ + --output /tmp/cube-experts --episodes 8 --record-video +``` + +`--episodes` 默认 8，范围 1–64；`--seed` 默认 13，`--cuda-device` 默认 0。输出目录必须不存在或为空。CPU 物理 200 Hz，控制及视频 20 Hz。采集版提高机械臂刚度到 200000，并将夹爪打开目标设在关节下限内侧 1 mm，以满足既定跟踪门槛并避免自由段触及限位；重力保持开启。 + +验收包括： + +- 规划使用 URDF 碰撞形状的保守凸包，按完整关节状态检查手指和 mimic 几何，并检查持物路径与桌面、地面、其他物体的碰撞。非相邻自碰撞也参与检查，结构上相距两条关节边以内的 link 对沿用现有 cuRobo 排除规则。 +- 实测关节及 cube 位姿再次检查几何，每个 5 ms 物理子步检查原生接触。只有声明的指尖、方块、支撑和固定安装部位接触被允许；接近段指尖接触限于 TCP 距方块 6 cm 的进入区域。未知物体、跨行接触、超过 2 mm 的穿透或接触数据超限都会拒绝。 +- 闭合后抬升并保持 1.5 秒，要求方块至少抬高 12 cm、两个手指各有至少 95% 的保持采样出现真实接触力，且从抬升到保持的 TCP 相对漂移不超过 1 cm / 0.15 rad。原有速度、加速度、轨迹质量和末端跟踪检查继续生效。 +- 只有验收通过、写入封存且成功回读的 episode 才进入 manifest。视频也会显示未通过的尝试，不能用视频代替数据验收。 + +输出 `preview.mp4`（开启视频时）、`generation_report.json`、`pickup_report.json`、`manifest.json` 和逐 episode 的 LeRobot 分片。每条数据含 T 个真实控制目标及 T+1 个实测观测/时间戳。cube/TCP 位姿在 LeRobot 中按行优先展开为 16 维，`episode.json` 的 `observation_shapes` 保留 `[4, 4]`；终态 `terminal.npz` 保持矩阵形状。`commanded_joint_indices` 标识完整关节向量中真正下发的关节列。 + +范围：单个固定基座 URDF 机器人、盒状刚体、CPU 物理。碰撞验证是有界采样，未提供连续碰撞保证。`source.kind=atomic` 表示来源是 `AtomicActionEngine.compile` 的离线计划；执行采用单独验收的 qpos 回放，不提交编译时预测的符号效果。Atomic Runtime 的 tracking/recovery 接入、Gym 接触采集与通用 YAML 启动器仍是后续工作。 + +## 自由运动采集 + +`free_motion.py` 使用普通重力下的纯机械臂 UR5，在固定场景中执行关节轨迹，经过规划、实测碰撞、任务和运动质量验收后保存一条 LeRobot episode。仿真物理在 CPU 上运行，cuRobo 和可选离屏相机使用 CUDA，无需桌面窗口。当前示例不包含抓取或持物。 + +在已安装 EmbodiChain 仿真依赖、CUDA、cuRobo 和 LeRobot 的环境中，从仓库根目录运行。保存视频还需要 `imageio-ffmpeg`： + +```bash +python -m pip install imageio-ffmpeg +python examples/sim/motion/trajectory_generation/free_motion.py --output /tmp/ur5-video --record-video --duration 5 --joint-displacement 0.4 +``` + +输出目录必须不存在或为空。上述命令记录 5 秒物理运动，使用 640×480 离屏相机流式写入 20 fps 的 H.264 视频。包含初始帧和终态帧，共 101 帧，因此视频时长为 5.05 秒。 + +不需要视频时可直接运行默认的 1 秒轨迹： + +```bash +python examples/sim/motion/trajectory_generation/free_motion.py --output /tmp/ur5-free-motion +``` + +可用参数： + +| 参数 | 含义 | +|---|---| +| `--record-video` | 在输出目录保存 `preview.mp4` | +| `--duration` | 运动时长，默认 1 秒；范围 0.05–30 秒，必须是 0.05 的整数倍 | +| `--joint-displacement` | 第一个机械臂关节的正向位移，默认 0.08 rad；范围 `(0, 0.5]` rad | +| `--cuda-device` | CUDA 设备编号，默认 0 | +| `--robot panda` | 运行 Panda 负例；当前锁定夹爪模型会拒绝普通重力下的实测手指漂移 | + +输出文件： + +- `preview.mp4`：启用视频时的实际执行记录，独立于数值 LeRobot 观测。 +- `generation_report.json`：计数、验收证据、失败原因及 `target_reached`。 +- `manifest.json`：已确认提交的 episode 分片目录；被拒绝的轨迹不会作为已提交数据列入。 +- 每个提交分片中的 `dataset/`、`terminal.npz` 和 `episode.json`：LeRobot 训练帧、终态/T+1 时间证据和谱系/验证信息。 + +视频存在不代表任务或专家数据验收通过，应查看报告和 manifest。未达到采集目标时脚本返回退出码 1；改变时长或位移后也会执行全部验收检查。 diff --git a/examples/sim/motion/trajectory_generation/cube_grasp_parallel.py b/examples/sim/motion/trajectory_generation/cube_grasp_parallel.py new file mode 100644 index 000000000..17d572d64 --- /dev/null +++ b/examples/sim/motion/trajectory_generation/cube_grasp_parallel.py @@ -0,0 +1,544 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Show grasp-pose and path augmentation on four synchronized cube pickups. + +Run from the repository root with the simulation and video dependencies:: + + python examples/sim/motion/trajectory_generation/cube_grasp_parallel.py --output /tmp/cube-grasp-parallel + +The four real physics rows share the same initial robot and cube poses. A cube +symmetry supplies two grasp orientations; joint residuals supply two transit +paths per orientation. Existing Atomic Skills plan the final approach, close, +and lift. The cube moves only through physical contact with the gripper. + +This example saves a four-panel MP4, measured rollout arrays, and a grasp-result +report. It does not use the free-motion-only GenerationRunner or certify a +contact-aware expert dataset. Planning uses the existing UR IK solver and the +MotionGenerator ik_interp strategy, without a collision-planning guarantee. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +import sys + +import imageio.v2 as imageio +import numpy as np +from PIL import Image, ImageDraw, ImageFont +import torch + +_REPO_ROOT = Path(__file__).resolve().parents[4] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.atomic_actions import ( + ControlPartCommandProfile, + EndEffectorPoseGoal, + GraspGoal, + MotionPolicy, + PickUpOptions, + create_simulation_atomic_action_engine, +) +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg +from embodichain.lab.sim.material import VisualMaterialCfg +from embodichain.lab.sim.motion.trajectory_augmentation import ( + TrajectoryPhase, + TrajectoryTemplate, + joint_residual, + rotate_grasp_about_object_axis, + validate_motion_limits, +) +from embodichain.lab.sim.objects import RigidObjectCfg +from embodichain.lab.sim.sensors import CameraCfg +from embodichain.lab.sim.shapes import CubeCfg +from embodichain.utils.math import look_at_to_pose +from scripts.tutorials.atomic_action.tutorial_utils import ( + create_antipodal_semantics, + create_toppra_motion_generator, + create_ur5_gripper_robot_cfg, + get_hand_open_close_qpos, +) + +__all__ = ["run_cube_grasp_parallel", "main"] + +_FPS = 20 +_CONTROL_DT = 1 / _FPS +_PHYSICS_DT = 0.005 +_WIDTH, _HEIGHT = 640, 480 +_HEADER = 48 +_YAW = (0.0, 0.0, math.pi / 2, math.pi / 2) +_COLORS = ((79, 183, 244), (250, 180, 63), (177, 143, 245), (80, 213, 160)) +_EYE, _TARGET = (-1.0, -1.0, 1.12), (-0.28, -0.08, 0.56) +_INTRINSICS = (570.0, 570.0, 320.0, 240.0) + + +def _grasp_results( + cube_poses: np.ndarray, tcp_poses: np.ndarray, hold_start: int +) -> list[dict[str, object]]: + """Check sustained lift and object stability relative to the actual TCP.""" + results = [] + for row in range(4): + cube_hold = cube_poses[row, hold_start:] + tcp_hold = tcp_poses[row, hold_start:] + relative = np.linalg.inv(tcp_hold) @ cube_hold + min_lift = float((cube_hold[:, 2, 3] - cube_poses[row, 0, 2, 3]).min()) + drift = float( + np.linalg.norm(relative[:, :3, 3] - relative[0, :3, 3], axis=-1).max() + ) + distance = float(np.linalg.norm(relative[:, :3, 3], axis=-1).max()) + results.append( + { + "row": row, + "yaw_degrees": round(math.degrees(_YAW[row])), + "path": "residual" if row % 2 else "reference", + "success": min_lift >= 0.12 and drift <= 0.01 and distance <= 0.06, + "min_hold_lift_m": min_lift, + "hold_relative_drift_m": drift, + "max_cube_tcp_distance_m": distance, + } + ) + return results + + +def _preview_frame( + rgb: np.ndarray, + trails: list[list[tuple[float, float]]], + elapsed: float, + phase: str, + lifts: np.ndarray, + font: ImageFont.ImageFont, + path_labels: tuple[str, ...] | None = None, +) -> np.ndarray: + """Tile camera frames from the same simulation tick with measured TCP trails.""" + preview = Image.new("RGB", (2 * _WIDTH, 2 * (_HEIGHT + _HEADER)), (18, 23, 32)) + for row in range(4): + panel = Image.fromarray(rgb[row]) + draw = ImageDraw.Draw(panel) + if len(trails[row]) > 1: + draw.line(trails[row], fill=_COLORS[row], width=3) + left, top = (row % 2) * _WIDTH, (row // 2) * (_HEIGHT + _HEADER) + preview.paste(panel, (left, top + _HEADER)) + draw = ImageDraw.Draw(preview) + label = ( + path_labels[row] + if path_labels is not None + else ("reference path" if row % 2 == 0 else "augmented path") + ) + draw.text( + (left + 12, top + 3), + f"{row + 1} | grasp {round(math.degrees(_YAW[row]))} deg | {label}", + fill=_COLORS[row], + font=font, + ) + draw.text( + (left + 12, top + 25), + f"{elapsed:5.2f} s | {phase} | cube lift {lifts[row] * 100:5.1f} cm", + fill=(229, 234, 241), + font=font, + ) + return np.asarray(preview) + + +def _prepare_cube_scene( + sim: SimulationManager, *, arm_stiffness: float = 5e4, hand_open_margin: float = 0.0 +): + """Share the four-row physical scene and initial state between examples.""" + robot_cfg = create_ur5_gripper_robot_cfg(init_pos=(0.0, 0.0, 0.3), tcp_z=0.15) + robot_cfg.drive_pros.stiffness["arm"] = arm_stiffness + robot = sim.add_robot(cfg=robot_cfg) + sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="bench", + body_type="kinematic", + init_pos=(-0.2, -0.08, 0.15), + shape=CubeCfg( + size=(0.9, 0.65, 0.3), + visual_material=VisualMaterialCfg( + uid="bench_material", base_color=[0.16, 0.20, 0.27, 1.0] + ), + ), + ) + ) + cube = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="cube", + shape=CubeCfg( + size=(0.05, 0.05, 0.05), + visual_material=VisualMaterialCfg( + uid="cube_material", base_color=[0.95, 0.23, 0.04, 1.0] + ), + ), + init_pos=(-0.42, -0.08, 0.325), + attrs=RigidBodyAttributesCfg( + mass=0.05, dynamic_friction=0.97, static_friction=0.99 + ), + ) + ) + camera = sim.add_sensor( + CameraCfg( + uid="comparison_camera", + visualization_role="record", + width=_WIDTH, + height=_HEIGHT, + far=2.5, + intrinsics=_INTRINSICS, + extrinsics=CameraCfg.ExtrinsicsCfg(eye=_EYE, target=_TARGET), + ) + ) + # Settle the identical cubes before fixing the reference state. No + # object pose, gravity, or dynamics are changed during the rollout. + sim.update(step=100) + object_poses = cube.get_local_pose(to_matrix=True).clone() + torch.testing.assert_close( + object_poses, object_poses[0:1].expand_as(object_poses), atol=1e-5, rtol=0 + ) + reference_grasp = torch.eye(4, device=robot.device) + reference_grasp[:3, :3] = reference_grasp.new_tensor( + [[0.0, -1.0, 0.0], [-1.0, 0.0, 0.0], [0.0, 0.0, -1.0]] + ) + reference_grasp[:3, 3] = object_poses[0, :3, 3] + grasps = rotate_grasp_about_object_axis( + object_poses[0], + reference_grasp, + axis=torch.tensor([0.0, 0.0, 1.0]), + angles=torch.tensor(_YAW), + ) + start_pose = reference_grasp.repeat(4, 1, 1) + start_pose[:, :3, 3] += start_pose.new_tensor([0.12, 0.0, 0.34]) + success, arm_qpos = robot.compute_ik( + pose=start_pose, joint_seed=robot.get_qpos(name="arm"), name="arm" + ) + if not success.all(): + raise RuntimeError("The common initial TCP pose is unreachable") + hand_open, hand_close = get_hand_open_close_qpos(robot) + hand_open = hand_open + hand_open_margin + initial_qpos = torch.tensor(robot.cfg.init_qpos, device=robot.device).repeat(4, 1) + initial_qpos[:, robot.get_joint_ids("arm")] = arm_qpos + initial_qpos[:, robot.get_joint_ids("hand")] = hand_open + for child, parent, multiplier, offset in zip( + robot.mimic_ids, + robot.mimic_parents, + robot.mimic_multipliers, + robot.mimic_offsets, + ): + initial_qpos[:, child] = initial_qpos[:, parent] * multiplier + offset + for target in (False, True): + robot.set_qpos(initial_qpos, target=target) + robot.clear_dynamics() + + return robot, cube, camera, initial_qpos, grasps, hand_open, hand_close + + +def _compile_cube_pickup(robot, cube, grasps, hand_open, hand_close): + """Compile the same atomic transit and PickUp for preview and collection.""" + generator = create_toppra_motion_generator(robot) + engine = create_simulation_atomic_action_engine( + motion_generator=generator, + scene_entities=(cube,), + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, grasp=hand_close + ) + }, + ) + pregrasp = grasps.clone() + pregrasp[:, 2, 3] += 0.16 + compiled = engine.compile( + ( + engine.make_invocation( + "move_end_effector", + EndEffectorPoseGoal(pregrasp), + control_parts={"primary": {"motion": "arm"}}, + motion_policy=MotionPolicy(strategy="ik_interp", sample_count=81), + ), + engine.make_invocation( + "pick_up", + GraspGoal( + create_antipodal_semantics(cube, label="cube"), + grasp_xpos=grasps, + ), + control_parts={"primary": {"motion": "arm", "grasp": "hand"}}, + motion_policy=MotionPolicy(strategy="ik_interp", sample_count=141), + skill_options=PickUpOptions( + pre_grasp_distance=0.16, + lift_height=0.18, + hand_interp_steps=30, + grasp_settle_steps=20, + ), + ), + ), + engine.initial_context(control_dt=_CONTROL_DT), + ) + return generator, compiled + + +def run_cube_grasp_parallel( + output_dir: Path, *, seed: int = 13, cuda_device: int = 0 +) -> dict[str, object]: + """Run four simultaneous physical pickups and save their measured evidence. + + Args: + output_dir: New or empty directory for the MP4, arrays, and report. + seed: Local random seed for both residual variants; no global RNG reset. + cuda_device: GPU used by the offscreen renderer. + + Returns: + The persisted report with per-row grasp outcomes and augmentation factors. + + Raises: + ValueError: If the output directory is occupied or the seed is invalid. + RuntimeError: If initialization, IK, or sampled motion-limit checks fail. + """ + output_dir = Path(output_dir) + if output_dir.exists() and (not output_dir.is_dir() or any(output_dir.iterdir())): + raise ValueError("output_dir must be new or empty") + if type(seed) is not int or not 0 <= seed < 2**63: + raise ValueError("seed must be an integer in [0, 2**63)") + output_dir.mkdir(parents=True, exist_ok=True) + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + sim_device="cpu", + num_envs=4, + physics_dt=_PHYSICS_DT, + arena_space=4.0, + gpu_id=cuda_device, + ) + ) + generator = writer = None + try: + robot, cube, camera, initial_qpos, grasps, hand_open, hand_close = ( + _prepare_cube_scene(sim) + ) + generator, compiled = _compile_cube_pickup( + robot, cube, grasps, hand_open, hand_close + ) + if not compiled.plan_success.all(): + raise RuntimeError( + f"IK planning failed in rows: {compiled.plan_success.tolist()}" + ) + transit_stop = compiled.action_waypoint_offset(1) + positions = compiled.trajectory.positions.clone() + # Atomic concatenation retains the next action's zero-time start. + # Here that duplicate becomes one explicit control-period hold so all + # four rows and their camera frames share a uniform execution clock. + torch.testing.assert_close( + positions[:, transit_stop], + positions[:, transit_stop - 1], + atol=1e-5, + rtol=0, + ) + intervals = compiled.trajectory.dt.clone() + intervals[:, transit_stop] = _CONTROL_DT + baseline = positions.clone() + arm_ids = tuple(robot.get_joint_ids("arm")) + limits = robot.get_qpos_limits()[0] + motion_checks = [] + for row in range(4): + template = TrajectoryTemplate( + "cube_pickup", + "v1", + f"yaw_{round(math.degrees(_YAW[row]))}", + tuple(robot.joint_names), + positions[row], + intervals[row], + ( + TrajectoryPhase( + "transit", + 0, + transit_stop, + allowed_operators=("joint_residual",), + ), + TrajectoryPhase( + "pickup", transit_stop, positions.shape[1], kind="contact" + ), + ), + allowed_operators=("joint_residual",), + controlled_joint_indices=arm_ids, + ) + if row % 2: + template = joint_residual( + template, + joint_limits=limits, + normalized_scale=0.025, + generator=torch.Generator().manual_seed(seed), + ) + check = validate_motion_limits( + template, + velocity_limits=torch.full((robot.dof,), 1.5), + acceleration_limits=torch.full((robot.dof,), 15.0), + ) + if not check.accepted: + raise RuntimeError(f"Motion limits failed in row {row}: {check}") + motion_checks.append(dict(check.checks[0].metrics)) + positions[row] = template.positions + # The augmentation is confined to transit: all closing/contact/lift + # commands and the endpoint into the pre-grasp corridor are preserved. + assert torch.equal( + positions[:, transit_stop - 1 :], baseline[:, transit_stop - 1 :] + ) + torch.testing.assert_close(positions[:, 0], robot.get_qpos(), atol=1e-5, rtol=0) + count = positions.shape[1] + hold_frames = 30 + positions = torch.cat( + (positions, positions[:, -1:].repeat(1, hold_frames, 1)), dim=1 + ) + try: + font = ImageFont.truetype("DejaVuSans.ttf", 17) + except OSError: + font = ImageFont.load_default() + writer = imageio.get_writer( + str(output_dir / "preview.mp4"), + format="FFMPEG", + fps=_FPS, + codec="libx264", + pixelformat="yuv420p", + quality=8, + ffmpeg_log_level="error", + output_params=["-movflags", "+faststart"], + ) + view = torch.linalg.inv(look_at_to_pose(_EYE, _TARGET)).squeeze(0).numpy() + trails: list[list[tuple[float, float]]] = [[] for _ in range(4)] + qpos_samples, cube_samples, tcp_samples, timestamps = [], [], [], [] + start_time = sim.simulation_time + close_start = compiled.segment(1, "close").start + lift_start = compiled.segment(1, "lift").start + + for index in range(positions.shape[1]): + if index: + robot.set_qpos(positions[:, index]) + sim.update(step=round(_CONTROL_DT / _PHYSICS_DT)) + elapsed = sim.simulation_time - start_time + if not math.isclose(elapsed, index * _CONTROL_DT, abs_tol=1e-7): + raise RuntimeError( + "The measured simulation clock diverged from the video clock" + ) + qpos_samples.append(robot.get_qpos().cpu().numpy().copy()) + cube_samples.append( + cube.get_local_pose(to_matrix=True).cpu().numpy().copy() + ) + tcp = robot.compute_fk( + robot.get_qpos(name="arm"), name="arm", to_matrix=True + ) + tcp_samples.append(tcp.cpu().numpy().copy()) + timestamps.append(elapsed) + for row in range(4): + point = view @ np.r_[tcp_samples[-1][row, :3, 3], 1.0] + if point[2] > 0: + fx, fy, cx, cy = _INTRINSICS + trails[row].append( + (fx * point[0] / point[2] + cx, fy * point[1] / point[2] + cy) + ) + phase = ( + "transit" + if index < transit_stop + else ( + "approach" + if index < close_start + else ( + "close gripper" + if index < lift_start + else "lift" if index < count else "hold" + ) + ) + ) + camera.update() + rgb = camera.get_data()["color"][..., :3].cpu().numpy().copy() + writer.append_data( + _preview_frame( + rgb, + trails, + elapsed, + phase, + cube_samples[-1][:, 2, 3] - cube_samples[0][:, 2, 3], + font, + ) + ) + + cubes = np.stack(cube_samples, axis=1) + tcps = np.stack(tcp_samples, axis=1) + results = _grasp_results(cubes, tcps, count) + deviations = [ + float( + np.linalg.norm( + tcps[row, :transit_stop, :3, 3] + - tcps[row - 1, :transit_stop, :3, 3], + axis=-1, + ).max() + ) + for row in (1, 3) + ] + report = { + "success": all(result["success"] for result in results), + "num_envs": 4, + "seed": seed, + "normalized_residual_scale": 0.025, + "fps": _FPS, + "frame_count": len(timestamps), + "physical_duration_s": timestamps[-1], + "transit_stop": transit_stop, + "hold_start": count, + "contact_commands_preserved": True, + "max_path_pair_separation_m": deviations, + "planned_motion_limits": motion_checks, + "results": results, + "validation_scope": "physical sustained lift and TCP-relative position stability; no contact-aware collision or expert-dataset certification", + } + np.savez_compressed( + output_dir / "rollout.npz", + timestamps=np.asarray(timestamps), + commanded_qpos=positions.cpu().numpy(), + measured_qpos=np.stack(qpos_samples, axis=1), + cube_poses=cubes, + tcp_poses=tcps, + grasp_poses=grasps.cpu().numpy(), + reference_qpos=baseline.cpu().numpy(), + ) + (output_dir / "report.json").write_text(json.dumps(report, indent=2) + "\n") + return report + finally: + try: + if writer is not None: + writer.close() + finally: + if generator is not None: + generator.planner.close() + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + +def main() -> None: + """Run the parallel comparison and return a failing exit code for failed grasps.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--seed", type=int, default=13) + parser.add_argument("--cuda-device", type=int, default=0) + args = parser.parse_args() + report = run_cube_grasp_parallel( + args.output, seed=args.seed, cuda_device=args.cuda_device + ) + print(json.dumps(report)) + if not report["success"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/sim/motion/trajectory_generation/cube_pickup_collection.py b/examples/sim/motion/trajectory_generation/cube_pickup_collection.py new file mode 100644 index 000000000..1a66f08ad --- /dev/null +++ b/examples/sim/motion/trajectory_generation/cube_pickup_collection.py @@ -0,0 +1,416 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Collect contact-validated parallel cube PickUp episodes and an optional video. + +Run from the repository root:: + + python examples/sim/motion/trajectory_generation/cube_pickup_collection.py --output /tmp/cube-experts --episodes 8 --record-video + +Requires the simulator, LeRobot, python-fcl, trimesh and yourdfpy. CPU physics +is stepped at 200 Hz; a CUDA renderer is needed for the optional 20 Hz preview. +Four rows share one physical clock. The complete batch restores its initial +robot/object state between rounds. Only validated, sealed and read-back episodes +count towards the requested target. The video also shows rejected attempts. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import sys +import traceback + +import numpy as np +import torch + +_REPO_ROOT = Path(__file__).resolve().parents[4] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.motion.trajectory_augmentation import ( + SceneCase, + TrajectoryGenerationJobCfg, + ValidationCheck, + ValidationResult, +) +from embodichain.lab.trajectory_generation.execution import QposRolloutExecutor +from embodichain.lab.trajectory_generation.initial_state import ( + FixedSceneHost, + InitialStateProfile, +) +from embodichain.lab.trajectory_generation.integrations.sim import ( + SimInitialStateAdapter, +) +from embodichain.lab.trajectory_generation.integrations.atomic import ( + export_pickup_templates, +) +from embodichain.lab.trajectory_generation.integrations.contact import ( + PickUpContactProfile, + PickUpMotionValidator, +) +from embodichain.lab.trajectory_generation.runner import ( + GenerationRunner, + MotionLimitsProfile, +) +from embodichain.lab.trajectory_generation.sinks import LeRobotEpisodeSink +from embodichain.utils.math import look_at_to_pose +from examples.sim.motion.trajectory_generation.cube_grasp_parallel import ( + _prepare_cube_scene, + _compile_cube_pickup, + _preview_frame, + _YAW, + _EYE, + _TARGET, + _INTRINSICS, + _CONTROL_DT, + _PHYSICS_DT, +) + +__all__ = ["run_cube_pickup_collection", "main"] + + +def run_cube_pickup_collection( + output_dir: Path, + *, + episodes: int = 8, + seed: int = 13, + cuda_device: int = 0, + record_video: bool = False, +) -> dict[str, object]: + """Run a bounded four-row PickUp collection with real contact acceptance. + + Args: + output_dir: New or empty local collection directory. + episodes: Number of committed episodes requested, between 1 and 64. + seed: Local augmentation random seed. + cuda_device: GPU used by the simulator renderer. + record_video: Stream synchronized four-panel observations to preview.mp4. + + Returns: + The persisted generation audit with committed count and target outcome. + """ + if type(episodes) is not int or not 1 <= episodes <= 64: + raise ValueError("episodes must be an integer between 1 and 64") + output_dir = Path(output_dir) + if output_dir.exists() and (not output_dir.is_dir() or any(output_dir.iterdir())): + raise ValueError("output_dir must be new or empty") + cfg = TrajectoryGenerationJobCfg.from_mapping( + { + "source": {"kind": "atomic", "source_id": "atomic_pickup"}, + "augmentation": { + "seed": seed, + "factors": { + "spatial": { + "enabled": True, + "method": "joint_residual", + "joint_offset_scale": 0.025, + } + }, + }, + "collection": { + "target_committed_episodes": episodes, + "max_proposals": max(32, episodes * 8), + "max_rollout_attempts": max(16, episodes * 4), + "max_wall_time_s": 1200.0, + }, + "validation": {"path_length_ratio_max": 2.0}, + } + ) + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + sim_device="cpu", + num_envs=4, + physics_dt=_PHYSICS_DT, + arena_space=4.0, + gpu_id=cuda_device, + ) + ) + generator = host = sink = writer = None + try: + robot, cube, camera, initial, grasps, hand_open, hand_close = ( + _prepare_cube_scene(sim, arm_stiffness=2e5, hand_open_margin=0.001) + ) + generator, compiled = _compile_cube_pickup( + robot, cube, grasps, hand_open, hand_close + ) + templates = export_pickup_templates(compiled, robot, control_dt=_CONTROL_DT) + initial_objects = { + uid: sim.get_rigid_object(uid).get_local_pose(to_matrix=True).clone() + for uid in sim.get_rigid_object_uid_list() + } + zero = torch.zeros_like(initial) + prepares = 0 + + def prepare() -> None: + nonlocal prepares + prepares += 1 + robot.set_qpos(initial, target=False) + robot.set_qvel(zero, target=False) + robot.set_qpos(initial, target=True) + robot.set_qvel(zero, target=True) + robot.set_qf(zero) + for uid, pose in initial_objects.items(): + obj = sim.get_rigid_object(uid) + obj.set_local_pose(pose) + obj.clear_dynamics() + + def signature() -> str: + return json.dumps( + { + "robot": robot.cfg.to_dict(), + "scene": { + uid: sim.get_rigid_object(uid).cfg.to_dict() + for uid in initial_objects + }, + "physics": sim.sim_config.physics_config.to_dict(), + }, + sort_keys=True, + default=str, + ) + + def verify(cases) -> ValidationResult: + passed = torch.allclose( + robot.get_qpos(), initial, atol=1e-6, rtol=0 + ) and all( + torch.allclose( + sim.get_rigid_object(uid).get_local_pose(to_matrix=True), + pose, + atol=1e-6, + rtol=0, + ) + for uid, pose in initial_objects.items() + ) + return ValidationResult( + ( + ValidationCheck( + "task_initial_state", + "passed" if passed else "failed", + "Identical prepared cube, table and full robot state", + ), + ) + ) + + host = FixedSceneHost( + SimInitialStateAdapter(sim, robot), + InitialStateProfile( + profile_id="fixed_scene_initial_state", + prepare=prepare, + signature=signature, + verify=verify, + physics_dt=_PHYSICS_DT, + settling_steps=0, + ), + ) + # The PGI fingertips extend past the 0.15 m TCP. Their contact entry + # region includes the cube half-height and that physical overhang. + planner = PickUpMotionValidator( + sim, robot, profile=PickUpContactProfile(approach_contact_distance=0.06) + ) + frames = 0 + trails = [[] for _ in range(4)] + view = torch.linalg.inv(look_at_to_pose(_EYE, _TARGET)).squeeze(0).numpy() + if record_video: + from PIL import ImageFont + + try: + font = ImageFont.truetype("DejaVuSans.ttf", 17) + except OSError: + font = ImageFont.load_default() + + def observe() -> dict[str, torch.Tensor]: + nonlocal frames, trails + if record_video: + index = frames % len(templates[0].positions) + if index == 0: + trails = [[] for _ in range(4)] + observations = planner.observations() + tcp = observations["tcp_pose"].cpu().numpy() + for row in range(4): + point = view @ np.r_[tcp[row, :3, 3], 1.0] + if point[2] > 0: + fx, fy, cx, cy = _INTRINSICS + trails[row].append( + ( + fx * point[0] / point[2] + cx, + fy * point[1] / point[2] + cy, + ) + ) + phase = next( + p.phase_id + for p in templates[0].phases + if p.start_index <= index < p.stop_index + ) + camera.update() + rgb = camera.get_data()["color"][..., :3].cpu().numpy().copy() + frame = _preview_frame( + rgb, + trails, + index * _CONTROL_DT, + f"round {prepares}: {phase}", + ( + observations["object_pose"][:, 2, 3] + - initial_objects[cube.uid][:, 2, 3] + ) + .cpu() + .numpy(), + font, + path_labels=("augmented path",) * 4, + ) + writer.append_data(frame) + frames += 1 + return { + "joint_positions": robot.get_qpos(), + "joint_velocities": robot.get_qvel(), + } + + arm_ids = tuple(robot.get_joint_ids("arm")) + + def task_validator( + candidate, observations, actions, timestamps + ) -> ValidationResult: + measured = observations["joint_positions"] + expected = candidate.positions[0, : len(measured)].to(measured) + endpoint = float( + (measured[-1, arm_ids] - expected[-1, arm_ids]).abs().max() + ) + tracking = float((measured[:, arm_ids] - expected[:, arm_ids]).abs().max()) + # Physical grasp and contact checks are independent mandatory gates + # owned by planner; this check proves the commanded arm path executed. + return ValidationResult( + ( + ValidationCheck( + "task_success", + "passed" if endpoint <= 0.05 and tracking <= 0.08 else "failed", + "Measured arm endpoint and tracking", + { + "endpoint_error_rad": endpoint, + "tracking_error_rad": tracking, + }, + ), + ) + ) + + budget = 1024 * 1024 + executor = QposRolloutExecutor( + host, + control_dt=_CONTROL_DT, + observe=observe, + validator=task_validator, + max_episode_bytes=budget, + contact_validator=planner, + ) + sink = LeRobotEpisodeSink(output_dir, fps=20, max_episode_bytes=budget) + if record_video: + import imageio.v2 as imageio + + writer = imageio.get_writer( + str(output_dir / "preview.mp4"), + format="FFMPEG", + fps=20, + codec="libx264", + pixelformat="yuv420p", + quality=8, + ffmpeg_log_level="error", + output_params=["-movflags", "+faststart"], + ) + cases = tuple( + SceneCase( + f"cube_grasp_row_{row}", + f"fixed_yaw_{round(np.degrees(_YAW[row]))}", + "cube_table_v1", + "cube_pickup", + "ur5_dh_pgi", + ) + for row in range(4) + ) + limits = MotionLimitsProfile( + torch.full((robot.dof,), 1.5), torch.full((robot.dof,), 15.0) + ) + runner = GenerationRunner( + cfg, host, planner, executor, sink, motion_limits=limits + ) + report = dict(runner.run(cases, templates)) + (output_dir / "pickup_report.json").write_text( + json.dumps( + { + "committed": report["counts"]["committed"], + "target_reached": report["target_reached"], + "prepared_batches": prepares, + "video_frames": frames, + "source": "offline AtomicActionEngine MoveEndEffector + PickUp export", + "contact_profile": planner.profile.to_dict(), + }, + indent=2, + ) + ) + return report + finally: + if writer is not None: + writer.close() + if sink is not None: + sink.close() + if host is not None: + host.close() + if generator is not None: + generator.planner.close() + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + +def main() -> None: + """Run the collection CLI and exit unsuccessfully if its target was not met.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--episodes", type=int, default=8) + parser.add_argument("--seed", type=int, default=13) + parser.add_argument("--cuda-device", type=int, default=0) + parser.add_argument("--record-video", action="store_true") + args = parser.parse_args() + report = None + try: + report = run_cube_pickup_collection( + args.output, + episodes=args.episodes, + seed=args.seed, + cuda_device=args.cuda_device, + record_video=args.record_video, + ) + except Exception as error: + diagnostic = traceback.TracebackException.from_exception(error) + # Release completed rollout frames before interpreter shutdown; their + # locals can otherwise retain native scene objects through a traceback. + traceback.clear_frames(error.__traceback__) + print("".join(diagnostic.format()), file=sys.stderr, end="") + if report is None: + raise SystemExit(1) + print( + json.dumps( + { + "target_reached": report["target_reached"], + "counts": dict(report["counts"]), + } + ) + ) + if not report["target_reached"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/sim/motion/trajectory_generation/free_motion.py b/examples/sim/motion/trajectory_generation/free_motion.py new file mode 100644 index 000000000..254342371 --- /dev/null +++ b/examples/sim/motion/trajectory_generation/free_motion.py @@ -0,0 +1,445 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Collect one measured UR5 free-motion episode with real physics and cuRobo. + +Run from the repository root with CUDA, cuRobo and LeRobot installed:: + + python examples/sim/motion/trajectory_generation/free_motion.py --output /tmp/ur5-free-motion + +Save a five-second offscreen camera preview:: + + python examples/sim/motion/trajectory_generation/free_motion.py --output /tmp/ur5-video --record-video --duration 5 --joint-displacement 0.4 + +The pure-arm UR5 is fixed 0.5 m above the ground, without a gripper or held +object. A registered cuboid models the simulator's implicit ground collider. +Physics uses CPU; cuRobo and the optional offscreen camera use CUDA. Video is +streamed to preview.mp4 at 20 fps, including the initial and terminal observation. +No desktop window is required. This is a free-motion benchmark, not a PickUp task. +All expert gates remain enabled, and a rejected episode stays out of LeRobot. +Gravity remains enabled. ``--robot panda`` is a physical negative example: +its moving gripper/mimic coordinates violate the current locked-hand model, +so collection rejects the evidence and exits with status 1. +""" + +from __future__ import annotations + +import argparse +from collections.abc import Mapping +import json +import math +from pathlib import Path +import sys + +import torch + +_REPO_ROOT = Path(__file__).resolve().parents[4] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg +from embodichain.lab.sim.motion.planners import ( + CuroboPlannerCfg, + CuroboWorldCfg, + MotionGenCfg, + MotionGenerator, +) +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CandidateTrajectoryBatch, + SceneCase, + TrajectoryGenerationJobCfg, + TrajectoryPhase, + TrajectoryTemplate, + ValidationCheck, + ValidationResult, +) +from embodichain.lab.sim.objects import RigidObjectCfg +from embodichain.lab.sim.robots import FrankaPandaCfg, URRobotCfg +from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.sensors import CameraCfg +from embodichain.lab.trajectory_generation.execution import QposRolloutExecutor +from embodichain.lab.trajectory_generation.initial_state import ( + FixedSceneHost, + InitialStateProfile, +) +from embodichain.lab.trajectory_generation.integrations.planning import ( + EnvRowMotionPlanner, +) +from embodichain.lab.trajectory_generation.integrations.sim import ( + SimInitialStateAdapter, +) +from embodichain.lab.trajectory_generation.runner import ( + GenerationRunner, + MotionLimitsProfile, +) +from embodichain.lab.trajectory_generation.sinks import LeRobotEpisodeSink + +__all__ = ["run_free_motion", "main"] + + +def run_free_motion( + output_dir: Path, + *, + robot_type: str = "ur5", + cuda_device: int = 0, + record_video: bool = False, + duration_s: float = 1.0, + joint_displacement: float = 0.08, +) -> Mapping[str, object]: + """Run one bounded collection and return the persisted final audit. + + Args: + output_dir: Dedicated local LeRobot collection directory. + robot_type: ``ur5`` for the pure-arm benchmark; ``panda`` for the + rejection example with a moving articulated gripper. + cuda_device: GPU used by cuRobo and the headless simulator engine. + record_video: Stream actual RGB observations to ``preview.mp4`` in the + output directory. The numeric LeRobot observations retain their schema. + duration_s: Motion duration in seconds, a multiple of 0.05, up to 30. + joint_displacement: Positive first-arm-joint displacement in radians. + + Returns: + The Runner report, including committed counts and validation diagnostics. + """ + if robot_type not in {"ur5", "panda"}: + raise ValueError("robot_type must be ur5 or panda") + physics_dt, control_dt = 0.005, 0.05 + if not math.isfinite(duration_s) or not control_dt <= duration_s <= 30: + raise ValueError("duration_s must be between 0.05 and 30 seconds") + steps = round(duration_s / control_dt) + if not math.isclose(steps * control_dt, duration_s, rel_tol=0, abs_tol=1e-8): + raise ValueError("duration_s must be a multiple of the 0.05 s control period") + if not math.isfinite(joint_displacement) or not 0 < joint_displacement <= 0.5: + raise ValueError("joint_displacement must be in (0, 0.5] radians") + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + sim_device="cpu", + num_envs=1, + physics_dt=physics_dt, + gpu_id=cuda_device, + ) + ) + host = sink = generator = video_writer = None + try: + cfg_type = URRobotCfg if robot_type == "ur5" else FrankaPandaCfg + robot_values = { + "uid": f"generation_{robot_type}", + "robot_type": robot_type, + "init_pos": [0.0, 0.0, 0.5], + "enable_gravity": True, + } + if robot_type == "ur5": + robot_values["init_qpos"] = [0.0, -1.57, 1.57, -1.57, -1.57, 0.0] + robot = sim.add_robot(cfg=cfg_type.from_dict(robot_values)) + # SimulationManager creates an unregistered 1000 x 1000 x 100 m ground + # cuboid centered at z=-50.001. This explicit matching collider gives + # the trusted profile and cuRobo an owned representation of that floor. + floor = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="ground_proxy", + shape=CubeCfg(size=(1000.0, 1000.0, 100.0)), + attrs=RigidBodyAttributesCfg(), + body_type="kinematic", + init_pos=(0.0, 0.0, -50.001), + ) + ) + camera = ( + sim.add_sensor( + CameraCfg( + uid="preview_camera", + visualization_role="record", + width=640, + height=480, + intrinsics=(520.0, 520.0, 320.0, 240.0), + extrinsics=CameraCfg.ExtrinsicsCfg( + eye=(0.6, -1.3, 1.35), + target=(-0.2, 0.0, 1.0), + ), + ) + ) + if record_video + else None + ) + initial = torch.tensor([robot.cfg.init_qpos], device=robot.device) + zeros = torch.zeros_like(initial) + arm_ids = tuple(robot.get_joint_ids("arm")) + + def prepare() -> None: + robot.set_gravity(True) + robot.set_qpos(initial, target=False) + robot.set_qvel(zeros, target=False) + robot.set_qpos(initial, target=True) + robot.set_qvel(zeros, target=True) + robot.set_qf(zeros) + + def signature() -> str: + return json.dumps( + { + "robot": robot.cfg.to_dict(), + "physics": sim.sim_config.physics_config.to_dict(), + "floor_size": (1000.0, 1000.0, 100.0), + "floor_center": (0.0, 0.0, -50.001), + "camera": None if camera is None else camera.cfg.to_dict(), + "render": sim.sim_config.render_cfg.to_dict(), + }, + sort_keys=True, + default=str, + ) + + def verify(cases: tuple[SceneCase, ...]) -> ValidationResult: + passed = ( + len(cases) == 1 + and torch.allclose(robot.get_qpos(), initial, atol=1e-6, rtol=0) + and torch.allclose(robot.get_qvel(), zeros, atol=1e-6, rtol=0) + ) + return ValidationResult( + ( + ValidationCheck( + "task_initial_state", + "passed" if passed else "failed", + "Neutral full-joint position and zero initial joint velocity.", + ), + ) + ) + + host = FixedSceneHost( + SimInitialStateAdapter(sim, robot), + InitialStateProfile( + profile_id="fixed_scene_initial_state", + prepare=prepare, + signature=signature, + verify=verify, + physics_dt=physics_dt, + settling_steps=0, + ), + ) + generator = MotionGenerator( + MotionGenCfg( + planner_cfg=CuroboPlannerCfg( + robot_uid=robot.uid, + cuda_device=cuda_device, + use_cuda_graph=False, + world=CuroboWorldCfg( + rigid_objects={floor.uid: floor}, + dynamic_obstacle_names=[floor.uid], + obstacle_representation="cuboid", + ), + ) + ) + ) + planner = EnvRowMotionPlanner( + generator, + control_part="arm", + held_object_ids=(), + max_joint_step=0.01, + max_validation_samples=max(128, steps + 1), + ) + + def task_validator( + candidate: CandidateTrajectoryBatch, + observations: Mapping[str, torch.Tensor], + actions: torch.Tensor, + timestamps: torch.Tensor, + ) -> ValidationResult: + measured = observations["joint_positions"] + expected = candidate.positions[0, : int(candidate.valid_length[0])].to( + measured + ) + endpoint_error = float((measured[-1] - expected[-1]).abs().max()) + tracking_error = float((measured - expected).abs().max()) + movement = float(measured[-1, arm_ids[0]] - measured[0, arm_ids[0]]) + gripper_ids = [index for index in range(robot.dof) if index not in arm_ids] + gripper_drift = ( + float((measured[:, gripper_ids] - expected[:, gripper_ids]).abs().max()) + if gripper_ids + else 0.0 + ) + passed = ( + endpoint_error <= 0.01 + and tracking_error <= 0.02 + and movement >= 0.875 * joint_displacement + ) + return ValidationResult( + ( + ValidationCheck( + "task_success", + "passed" if passed else "failed", + "Actual endpoint, measured tracking, and observed arm displacement.", + { + "endpoint_error": endpoint_error, + "tracking_error": tracking_error, + "measured_displacement": movement, + "max_gripper_drift": gripper_drift, + }, + ), + ) + ) + + def observe() -> Mapping[str, torch.Tensor]: + if camera is not None: + camera.update() + rgb = ( + camera.get_data()["color"][0, ..., :3] + .cpu() + .contiguous() + .numpy() + .copy() + ) + video_writer.append_data(rgb) + return { + "joint_positions": robot.get_qpos(), + "joint_velocities": robot.get_qvel(), + } + + executor = QposRolloutExecutor( + host, + control_dt=control_dt, + observe=observe, + validator=task_validator, + max_episode_bytes=128 * 1024, + ) + sink = LeRobotEpisodeSink( + Path(output_dir), fps=20, max_episode_bytes=128 * 1024 + ) + if record_video: + import imageio.v2 as imageio + + video_writer = imageio.get_writer( + str(sink.root / "preview.mp4"), + format="FFMPEG", + fps=20, + codec="libx264", + pixelformat="yuv420p", + quality=8, + ffmpeg_log_level="error", + output_params=["-movflags", "+faststart"], + ) + limits = MotionLimitsProfile( + torch.tensor( + [1.0 if index in arm_ids else 0.1 for index in range(robot.dof)] + ), + torch.tensor( + [5.0 if index in arm_ids else 1.0 for index in range(robot.dof)] + ), + ) + runner = GenerationRunner( + TrajectoryGenerationJobCfg.from_mapping( + { + "collection": { + "target_committed_episodes": 1, + "max_proposals": 1, + "max_rollout_attempts": 1, + "max_wall_time_s": 180.0, + } + } + ), + host, + planner, + executor, + sink, + motion_limits=limits, + ) + time = torch.linspace(0.0, 1.0, steps + 1) + smooth = 10 * time**3 - 15 * time**4 + 6 * time**5 + positions = initial.repeat(steps + 1, 1) + positions[:, arm_ids[0]] += joint_displacement * smooth + dt = torch.full((steps + 1,), control_dt, dtype=torch.float64) + dt[0] = 0 + template = TrajectoryTemplate( + "handwritten_qpos", + f"v1-duration={duration_s!r}-displacement={joint_displacement!r}", + "reference_0", + tuple(robot.joint_names), + positions, + dt, + (TrajectoryPhase("free", 0, steps + 1),), + validator_id="task_success", + controlled_joint_indices=arm_ids, + ) + case = SceneCase( + f"{robot_type}_free_motion", + "reference_start", + f"{robot_type}_elevated_ground_gravity_enabled", + "free_arm_motion", + robot.uid, + ) + return runner.run((case,), (template,)) + finally: + try: + if video_writer is not None: + video_writer.close() + finally: + if sink is not None: + sink.close() + if host is not None: + host.close() + if generator is not None: + generator.planner.close() + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + +def main() -> None: + """Parse the output location and run the measured free-motion benchmark.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--robot", choices=("ur5", "panda"), default="ur5") + parser.add_argument("--cuda-device", type=int, default=0) + parser.add_argument( + "--record-video", + action="store_true", + help="Save preview.mp4 using an offscreen RGB camera", + ) + parser.add_argument( + "--duration", + type=float, + default=1.0, + help="Motion duration in seconds (multiple of 0.05, max 30)", + ) + parser.add_argument( + "--joint-displacement", + type=float, + default=0.08, + help="First arm joint displacement in radians (0 < value <= 0.5)", + ) + args = parser.parse_args() + report = run_free_motion( + args.output, + robot_type=args.robot, + cuda_device=args.cuda_device, + record_video=args.record_video, + duration_s=args.duration, + joint_displacement=args.joint_displacement, + ) + video_path = args.output / "preview.mp4" + print( + json.dumps( + { + "counts": dict(report["counts"]), + "target_reached": report["target_reached"], + "report": str(args.output / "generation_report.json"), + "video": str(video_path) if video_path.is_file() else None, + } + ) + ) + if not report["target_reached"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/sim/workspace/analyze_cartesian_workspace.py b/examples/sim/motion/workspace/analyze_cartesian_workspace.py similarity index 97% rename from examples/sim/workspace/analyze_cartesian_workspace.py rename to examples/sim/motion/workspace/analyze_cartesian_workspace.py index fb9160067..4d2afd3da 100644 --- a/examples/sim/workspace/analyze_cartesian_workspace.py +++ b/examples/sim/motion/workspace/analyze_cartesian_workspace.py @@ -26,12 +26,12 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import MarkerCfg from embodichain.lab.sim.robots import DexforceW1Cfg -from embodichain.lab.sim.workspace import ( +from embodichain.lab.sim.motion.workspace import ( AnalysisMode, WorkspaceAnalyzer, WorkspaceAnalyzerConfig, ) -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( SamplingConfig, SamplingStrategy, VisualizationConfig, diff --git a/examples/sim/workspace/analyze_joint_workspace.py b/examples/sim/motion/workspace/analyze_joint_workspace.py similarity index 97% rename from examples/sim/workspace/analyze_joint_workspace.py rename to examples/sim/motion/workspace/analyze_joint_workspace.py index 3695bdb79..6b499f22a 100644 --- a/examples/sim/workspace/analyze_joint_workspace.py +++ b/examples/sim/motion/workspace/analyze_joint_workspace.py @@ -25,11 +25,11 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.robots import DexforceW1Cfg -from embodichain.lab.sim.workspace import ( +from embodichain.lab.sim.motion.workspace import ( WorkspaceAnalyzer, WorkspaceAnalyzerConfig, ) -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( SamplingConfig, SamplingStrategy, ) diff --git a/examples/sim/workspace/analyze_plane_workspace.py b/examples/sim/motion/workspace/analyze_plane_workspace.py similarity index 98% rename from examples/sim/workspace/analyze_plane_workspace.py rename to examples/sim/motion/workspace/analyze_plane_workspace.py index 95e381e1e..e8befba62 100644 --- a/examples/sim/workspace/analyze_plane_workspace.py +++ b/examples/sim/motion/workspace/analyze_plane_workspace.py @@ -26,12 +26,12 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import MarkerCfg from embodichain.lab.sim.robots import DexforceW1Cfg -from embodichain.lab.sim.workspace import ( +from embodichain.lab.sim.motion.workspace import ( AnalysisMode, WorkspaceAnalyzer, WorkspaceAnalyzerConfig, ) -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( SamplingConfig, SamplingStrategy, VisualizationConfig, diff --git a/pyproject.toml b/pyproject.toml index 299eaf49c..c59a28778 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,11 @@ dependencies = [ ] [project.optional-dependencies] +trajectory-generation = [ + "python-fcl>=0.7,<0.8", + "trimesh>=4.12,<5", + "yourdfpy>=0.0.60,<0.1", +] gensim = [ "bpy", "gradio>=6.17.3,<6.18", diff --git a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py index 7630f9641..bcd64fe03 100644 --- a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py +++ b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py @@ -237,8 +237,8 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, ) - from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg - from embodichain.lab.sim.planners import ToppraPlannerCfg + from embodichain.lab.sim.motion.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.motion.planners import ToppraPlannerCfg from scripts.tutorials.atomic_action.move_end_effector import ( create_robot, initialize_simulation, diff --git a/scripts/benchmark/atomic_action/move_held_object_benchmark.py b/scripts/benchmark/atomic_action/move_held_object_benchmark.py index 9e79cb374..34079aea3 100644 --- a/scripts/benchmark/atomic_action/move_held_object_benchmark.py +++ b/scripts/benchmark/atomic_action/move_held_object_benchmark.py @@ -657,8 +657,8 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: ensure_repo_root() ensure_torch() - from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg - from embodichain.lab.sim.planners import ToppraPlannerCfg + from embodichain.lab.sim.motion.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.motion.planners import ToppraPlannerCfg from scripts.tutorials.atomic_action.move_held_object import ( create_robot, initialize_simulation, diff --git a/scripts/benchmark/atomic_action/move_joints_benchmark.py b/scripts/benchmark/atomic_action/move_joints_benchmark.py index abdfcc6bd..3884c9203 100644 --- a/scripts/benchmark/atomic_action/move_joints_benchmark.py +++ b/scripts/benchmark/atomic_action/move_joints_benchmark.py @@ -244,8 +244,8 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: AtomicActionEngine, ControlPartCommandProfile, ) - from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg - from embodichain.lab.sim.planners import ToppraPlannerCfg + from embodichain.lab.sim.motion.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.motion.planners import ToppraPlannerCfg from scripts.tutorials.atomic_action.move_joints import ( create_robot, initialize_simulation, diff --git a/scripts/benchmark/atomic_action/pickup_benchmark.py b/scripts/benchmark/atomic_action/pickup_benchmark.py index 9fd4fea04..941252386 100644 --- a/scripts/benchmark/atomic_action/pickup_benchmark.py +++ b/scripts/benchmark/atomic_action/pickup_benchmark.py @@ -470,8 +470,8 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: ensure_repo_root() ensure_torch() - from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg - from embodichain.lab.sim.planners import ToppraPlannerCfg + from embodichain.lab.sim.motion.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.motion.planners import ToppraPlannerCfg from scripts.tutorials.atomic_action.pickup import ( create_robot, initialize_simulation, diff --git a/scripts/benchmark/atomic_action/place_benchmark.py b/scripts/benchmark/atomic_action/place_benchmark.py index 75c3e52d2..c0eabd3bb 100644 --- a/scripts/benchmark/atomic_action/place_benchmark.py +++ b/scripts/benchmark/atomic_action/place_benchmark.py @@ -637,8 +637,8 @@ def run_all_benchmarks(args: argparse.Namespace | None = None) -> Path: ensure_repo_root() ensure_torch() - from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg - from embodichain.lab.sim.planners import ToppraPlannerCfg + from embodichain.lab.sim.motion.planners import MotionGenerator, MotionGenCfg + from embodichain.lab.sim.motion.planners import ToppraPlannerCfg from scripts.tutorials.atomic_action.place import ( create_robot, initialize_simulation, diff --git a/scripts/benchmark/curobo_extraction/run_benchmark.py b/scripts/benchmark/curobo_extraction/run_benchmark.py index 1cd557ea5..5754d4938 100644 --- a/scripts/benchmark/curobo_extraction/run_benchmark.py +++ b/scripts/benchmark/curobo_extraction/run_benchmark.py @@ -44,12 +44,12 @@ import psutil import torch -from embodichain.lab.sim.planners.curobo.curobo_planner import ( +from embodichain.lab.sim.motion.planners.curobo.curobo_planner import ( CuroboPlanner, _CuroboBackend, _CuroboProfile, ) -from embodichain.lab.sim.planners.utils import PlanResult +from embodichain.lab.sim.motion.planners.utils import PlanResult # ============================================================================= # OLD (HEAD) implementations - verbatim loop logic, parameterized by device. diff --git a/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md b/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md index ae1a1bddc..348248bfb 100644 --- a/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md +++ b/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md @@ -50,7 +50,7 @@ physical task success. ### 2.1 Current NeuralPlanner capability boundary -`embodichain/lab/sim/planners/neural_planner.py` currently: +`embodichain/lab/sim/motion/planners/neural_planner.py` currently: - directly supports only `MoveType.EEF_MOVE`; - uses a 7-DoF waypoint Transformer checkpoint, currently centered on Franka; diff --git a/scripts/benchmark/motion_generation/metrics/trajectory.py b/scripts/benchmark/motion_generation/metrics/trajectory.py index fff97a6f1..f2bf9044c 100644 --- a/scripts/benchmark/motion_generation/metrics/trajectory.py +++ b/scripts/benchmark/motion_generation/metrics/trajectory.py @@ -24,7 +24,7 @@ import torch -from embodichain.lab.sim.planners.utils import PlanResult +from embodichain.lab.sim.motion.planners.utils import PlanResult from ..models import BenchmarkCase, CaseOutcome from .stats import nearest_rank_percentile diff --git a/scripts/benchmark/motion_generation/planners/base.py b/scripts/benchmark/motion_generation/planners/base.py index 895e5f11b..aeb0687e4 100644 --- a/scripts/benchmark/motion_generation/planners/base.py +++ b/scripts/benchmark/motion_generation/planners/base.py @@ -24,7 +24,7 @@ import torch -from embodichain.lab.sim.planners.utils import PlanResult +from embodichain.lab.sim.motion.planners.utils import PlanResult from ..config import PlannerSpecCfg, stable_hash from ..models import AlgorithmRole, BenchmarkCase, PlannerMetadata diff --git a/scripts/benchmark/motion_generation/planners/curobo.py b/scripts/benchmark/motion_generation/planners/curobo.py index 2b4e34716..09dcb8058 100644 --- a/scripts/benchmark/motion_generation/planners/curobo.py +++ b/scripts/benchmark/motion_generation/planners/curobo.py @@ -22,7 +22,7 @@ import torch -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( CuroboAutoGenCfg, CuroboPlanOptions, CuroboPlannerCfg, diff --git a/scripts/benchmark/motion_generation/planners/ik_interpolate.py b/scripts/benchmark/motion_generation/planners/ik_interpolate.py index f9ef082b1..1bbc77e14 100644 --- a/scripts/benchmark/motion_generation/planners/ik_interpolate.py +++ b/scripts/benchmark/motion_generation/planners/ik_interpolate.py @@ -22,7 +22,7 @@ import torch -from embodichain.lab.sim.planners import PlanResult +from embodichain.lab.sim.motion.planners import PlanResult from embodichain.lab.sim.utility.action_utils import interpolate_with_distance from ..models import BenchmarkCase diff --git a/scripts/benchmark/motion_generation/planners/neural.py b/scripts/benchmark/motion_generation/planners/neural.py index 20af60396..d54c9feb3 100644 --- a/scripts/benchmark/motion_generation/planners/neural.py +++ b/scripts/benchmark/motion_generation/planners/neural.py @@ -18,7 +18,7 @@ from __future__ import annotations -from embodichain.lab.sim.planners import PlanResult +from embodichain.lab.sim.motion.planners import PlanResult from ..models import BenchmarkCase from ..registry import register_planner_adapter diff --git a/scripts/benchmark/motion_generation/planners/toppra.py b/scripts/benchmark/motion_generation/planners/toppra.py index 038dbc0a7..4cdcb0ead 100644 --- a/scripts/benchmark/motion_generation/planners/toppra.py +++ b/scripts/benchmark/motion_generation/planners/toppra.py @@ -20,7 +20,7 @@ import importlib.util -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( MotionGenCfg, MotionGenOptions, MotionGenerator, @@ -29,7 +29,7 @@ ToppraPlannerCfg, ToppraPlanOptions, ) -from embodichain.lab.sim.planners.utils import TrajectorySampleMethod +from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod from ..config import PlannerSpecCfg from ..models import BenchmarkCase diff --git a/scripts/benchmark/motion_generation/runner.py b/scripts/benchmark/motion_generation/runner.py index 6f6badc87..3f39e7a56 100644 --- a/scripts/benchmark/motion_generation/runner.py +++ b/scripts/benchmark/motion_generation/runner.py @@ -25,7 +25,7 @@ import torch from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.planners.utils import PlanResult +from embodichain.lab.sim.motion.planners.utils import PlanResult from embodichain.lab.sim.robots import FrankaPandaCfg from . import planners as _builtin_planners # noqa: F401 - registry side effects diff --git a/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py b/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py index 88afa9bbf..553b50f85 100644 --- a/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py +++ b/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py @@ -35,9 +35,12 @@ import torch from embodichain.data import get_data_path -from embodichain.lab.sim.solvers.opw_solver import OPWSolverCfg -from embodichain.lab.sim.solvers.pytorch_solver import PytorchSolver, PytorchSolverCfg -from embodichain.lab.sim.solvers.ur_solver import URSolver, URSolverCfg +from embodichain.lab.sim.motion.solvers.opw_solver import OPWSolverCfg +from embodichain.lab.sim.motion.solvers.pytorch_solver import ( + PytorchSolver, + PytorchSolverCfg, +) +from embodichain.lab.sim.motion.solvers.ur_solver import URSolver, URSolverCfg OPW_LOWER_LIMITS = [-2.618, 0.0, -2.967, -1.745, -1.22, -2.0944] OPW_UPPER_LIMITS = [2.618, 3.14159, 0.0, 1.745, 1.22, 2.0944] diff --git a/scripts/benchmark/robotics/kinematic_solver/srs_solver.py b/scripts/benchmark/robotics/kinematic_solver/srs_solver.py index 4d8bee322..87e4f49b9 100644 --- a/scripts/benchmark/robotics/kinematic_solver/srs_solver.py +++ b/scripts/benchmark/robotics/kinematic_solver/srs_solver.py @@ -42,7 +42,7 @@ DexforceW1ArmSide, DexforceW1Version, ) -from embodichain.lab.sim.solvers.srs_solver import SRSSolverCfg +from embodichain.lab.sim.motion.solvers.srs_solver import SRSSolverCfg from embodichain.utils.logger import set_log_level DEFAULT_SIZES = (1, 16, 128) diff --git a/scripts/benchmark/workspace_analyzer/benchmark_workspace_analyzer.py b/scripts/benchmark/workspace_analyzer/benchmark_workspace_analyzer.py index ded6461b5..1421e7909 100644 --- a/scripts/benchmark/workspace_analyzer/benchmark_workspace_analyzer.py +++ b/scripts/benchmark/workspace_analyzer/benchmark_workspace_analyzer.py @@ -151,7 +151,7 @@ def benchmark_halton_sampler() -> ( tuple[list[dict[str, object]], list[dict[str, object]]] ): """Benchmark Halton sampler: vectorized vs loop-based.""" - from embodichain.lab.sim.workspace.samplers.halton_sampler import ( + from embodichain.lab.sim.motion.workspace.samplers.halton_sampler import ( HaltonSampler, ) @@ -215,10 +215,10 @@ def benchmark_density_metric() -> ( tuple[list[dict[str, object]], list[dict[str, object]]] ): """Benchmark density metric: KDTree vs brute-force.""" - from embodichain.lab.sim.workspace.metrics.density_metric import ( + from embodichain.lab.sim.motion.workspace.metrics.density_metric import ( DensityMetric, ) - from embodichain.lab.sim.workspace.configs.metric_config import ( + from embodichain.lab.sim.motion.workspace.configs.metric_config import ( DensityConfig, ) @@ -272,10 +272,10 @@ def benchmark_density_metric() -> ( def benchmark_voxelization() -> tuple[list[dict[str, object]], list[dict[str, object]]]: """Benchmark voxelization: np.unique vs dict-based.""" - from embodichain.lab.sim.workspace.metrics.reachability_metric import ( + from embodichain.lab.sim.motion.workspace.metrics.reachability_metric import ( ReachabilityMetric, ) - from embodichain.lab.sim.workspace.configs.metric_config import ( + from embodichain.lab.sim.motion.workspace.configs.metric_config import ( ReachabilityConfig, ) @@ -335,10 +335,10 @@ def benchmark_manipulability() -> ( tuple[list[dict[str, object]], list[dict[str, object]]] ): """Benchmark manipulability: batch vs per-sample.""" - from embodichain.lab.sim.workspace.metrics.manipulability_metric import ( + from embodichain.lab.sim.motion.workspace.metrics.manipulability_metric import ( ManipulabilityMetric, ) - from embodichain.lab.sim.workspace.configs.metric_config import ( + from embodichain.lab.sim.motion.workspace.configs.metric_config import ( ManipulabilityConfig, ) diff --git a/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py b/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py index c7a716a09..d5e0a0e3b 100644 --- a/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py +++ b/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py @@ -51,8 +51,8 @@ ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg from embodichain.lab.sim.objects import RigidObject, RigidObjectCfg, Robot -from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator -from embodichain.lab.sim.planners.curobo.curobo_planner import ( +from embodichain.lab.sim.motion.planners import MotionGenCfg, MotionGenerator +from embodichain.lab.sim.motion.planners.curobo.curobo_planner import ( CuroboAutoGenCfg, CuroboPlannerCfg, CuroboWorldCfg, diff --git a/scripts/tutorials/atomic_action/pour.py b/scripts/tutorials/atomic_action/pour.py index fced481d5..b749d8ef2 100644 --- a/scripts/tutorials/atomic_action/pour.py +++ b/scripts/tutorials/atomic_action/pour.py @@ -38,7 +38,10 @@ PourGoal, PourOptions, ) -from embodichain.lab.sim.planners import ToppraPlanOptions, TrajectorySampleMethod +from embodichain.lab.sim.motion.planners import ( + ToppraPlanOptions, + TrajectorySampleMethod, +) from embodichain.utils import logger from scripts.tutorials.atomic_action.axis_align import ( create_align_object, diff --git a/scripts/tutorials/atomic_action/scenario_utils.py b/scripts/tutorials/atomic_action/scenario_utils.py index a44c6ac52..8aaeb69e4 100644 --- a/scripts/tutorials/atomic_action/scenario_utils.py +++ b/scripts/tutorials/atomic_action/scenario_utils.py @@ -38,7 +38,7 @@ from embodichain.lab.sim.objects import RigidObject, Robot from embodichain.lab.sim.shapes import CubeCfg from embodichain.lab.sim.robots import build_dual_arm_cfg -from embodichain.lab.sim.solvers import PytorchSolverCfg +from embodichain.lab.sim.motion.solvers import PytorchSolverCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( ROBOTIQ_2F_140_TCP, diff --git a/scripts/tutorials/atomic_action/tutorial_utils.py b/scripts/tutorials/atomic_action/tutorial_utils.py index af8e288ce..185f5972e 100644 --- a/scripts/tutorials/atomic_action/tutorial_utils.py +++ b/scripts/tutorials/atomic_action/tutorial_utils.py @@ -38,7 +38,7 @@ ) from embodichain.lab.sim.cfg import LightCfg, MarkerCfg, RenderCfg, RobotCfg from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( CuroboPlannerCfg, MotionGenCfg, MotionGenerator, diff --git a/scripts/tutorials/grasp/grasp_generator.py b/scripts/tutorials/grasp/grasp_generator.py index afea548d1..30d87460b 100644 --- a/scripts/tutorials/grasp/grasp_generator.py +++ b/scripts/tutorials/grasp/grasp_generator.py @@ -31,7 +31,7 @@ from embodichain.lab.sim.objects import Robot, RigidObject from embodichain.lab.sim.utility.action_utils import interpolate_with_distance from embodichain.lab.sim.shapes import MeshCfg -from embodichain.lab.sim.solvers import URSolverCfg +from embodichain.lab.sim.motion.solvers import URSolverCfg from embodichain.data import get_data_path from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.toolkits.graspkit import ParallelJawGripperModelCfg diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index 98fd43e53..efa3c0f96 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -33,7 +33,7 @@ JointDrivePropertiesCfg, ) -from embodichain.lab.sim.solvers import PinkSolverCfg +from embodichain.lab.sim.motion.solvers import PinkSolverCfg from embodichain.data import get_data_path from embodichain.utils import logger diff --git a/scripts/tutorials/sim/motion_generator.py b/scripts/tutorials/sim/motion_generator.py index 776ca9dec..fc6619371 100644 --- a/scripts/tutorials/sim/motion_generator.py +++ b/scripts/tutorials/sim/motion_generator.py @@ -28,7 +28,7 @@ from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.cfg import RenderCfg from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( MotionGenCfg, MotionGenOptions, MotionGenerator, @@ -36,7 +36,7 @@ ToppraPlanOptions, ToppraPlannerCfg, ) -from embodichain.lab.sim.planners.utils import TrajectorySampleMethod +from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod from embodichain.lab.sim.robots import CobotMagicCfg RECORD_WIDTH = 1920 diff --git a/scripts/tutorials/sim/open_drawer.py b/scripts/tutorials/sim/open_drawer.py index 9e9e9ec3b..8ae5f8a55 100644 --- a/scripts/tutorials/sim/open_drawer.py +++ b/scripts/tutorials/sim/open_drawer.py @@ -33,7 +33,7 @@ RigidBodyAttributesCfg, ) from embodichain.lab.sim.objects import Articulation, Robot -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( MotionGenCfg, MotionGenerator, MotionGenOptions, diff --git a/scripts/tutorials/visualization/README.md b/scripts/tutorials/visualization/README.md index e809b7564..5533d7234 100644 --- a/scripts/tutorials/visualization/README.md +++ b/scripts/tutorials/visualization/README.md @@ -38,8 +38,8 @@ own `--host`/`--port` options: ```bash python scripts/tutorials/sim/create_scene.py --viser -python examples/sim/solvers/srs_solver.py --viser -python examples/sim/workspace/analyze_joint_workspace.py --viser +python examples/sim/motion/solvers/srs_solver.py --viser +python examples/sim/motion/workspace/analyze_joint_workspace.py --viser embodichain run-env --gym_config path/to/task.ur5.yaml --viser ``` diff --git a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py index d29a50af1..7ecfaf8ff 100644 --- a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py +++ b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py @@ -24,8 +24,8 @@ import pytest import torch -from embodichain.lab.sim.planners.curobo.curobo_planner import CuroboPlanner -from embodichain.lab.sim.planners.utils import MoveType, PlanResult +from embodichain.lab.sim.motion.planners.curobo.curobo_planner import CuroboPlanner +from embodichain.lab.sim.motion.planners.utils import MoveType, PlanResult from scripts.benchmark.motion_generation.aggregation import aggregate_results from scripts.benchmark.motion_generation.config import load_suite from scripts.benchmark.motion_generation.metrics.trajectory import ( diff --git a/tests/gym/envs/managers/test_workspace_randomization.py b/tests/gym/envs/managers/test_workspace_randomization.py index def266c2c..5cb96d185 100644 --- a/tests/gym/envs/managers/test_workspace_randomization.py +++ b/tests/gym/envs/managers/test_workspace_randomization.py @@ -23,7 +23,7 @@ from embodichain.lab.gym.envs.managers.randomization.spatial import ( sample_rigid_object_pose_from_workspace, ) -from embodichain.lab.sim.workspace import WorkspaceSample +from embodichain.lab.sim.motion.workspace import WorkspaceSample class _MockRobot: diff --git a/tests/gym/envs/task_program/test_simulation_environment.py b/tests/gym/envs/task_program/test_simulation_environment.py index 8cfc17c62..c0e410c24 100644 --- a/tests/gym/envs/task_program/test_simulation_environment.py +++ b/tests/gym/envs/task_program/test_simulation_environment.py @@ -85,7 +85,7 @@ RuntimeCommandFrame, RuntimeCommandPayload, ) -from embodichain.lab.sim.planners import MotionGenerator +from embodichain.lab.sim.motion.planners import MotionGenerator from embodichain.lab.task_program.semantics import ( BoundSemanticCall, COMPOSITE_EFFECT_MONITOR_ID, diff --git a/tests/gym/envs/task_program/test_simulation_parallel_safety.py b/tests/gym/envs/task_program/test_simulation_parallel_safety.py index 263ba5772..5e27f6846 100644 --- a/tests/gym/envs/task_program/test_simulation_parallel_safety.py +++ b/tests/gym/envs/task_program/test_simulation_parallel_safety.py @@ -34,7 +34,7 @@ JointPositionTarget, RuntimeCommandFrame, ) -from embodichain.lab.sim.planners import CuroboPlanner, MotionGenerator +from embodichain.lab.sim.motion.planners import CuroboPlanner, MotionGenerator from embodichain.lab.task_program.semantics import SceneRegistry from embodichain.lab.task_program.runtime.parallel_executor import ParallelSafetyError diff --git a/tests/gym/envs/test_demo.py b/tests/gym/envs/test_demo.py index 5b5c3fa5a..3413deaca 100644 --- a/tests/gym/envs/test_demo.py +++ b/tests/gym/envs/test_demo.py @@ -347,6 +347,125 @@ def test_handwritten_motion_generator_segment_declares_exact_progress_total() -> env._stack_block.clear_dynamics.assert_called_once() +@pytest.mark.parametrize("single_segment", [False, True]) +def test_execute_demo_episode_consumes_explicit_candidate_without_replanning( + single_segment: bool, +) -> None: + env = _SegmentedEnv() + env.create_demo_segments = Mock(side_effect=AssertionError("must not replan")) + env.metadata = {"dataset": {"instruction": {"lang": "Lift the selected object"}}} + env._begin_demo_episode_recording = Mock() + env._begin_demo_segment_recording = Mock() + env._end_demo_episode_recording = Mock() + candidate = DemoSegment(actions=(2, 3), name="candidate") + + result = execute_demo_episode( + env, + segments=candidate if single_segment else (candidate,), + episode_index=4, + attempt_id=2, + ) + + env.create_demo_segments.assert_not_called() + assert env.actions == [2, 3] + assert result.all_success + assert result.segments[0].name == "candidate" + assert result.segments[0].instruction == "Lift the selected object" + assert candidate.instruction is None + assert all(env.no_auto_reset_during_steps) + assert not env._demo_no_auto_reset + env._begin_demo_episode_recording.assert_called_once() + env._begin_demo_segment_recording.assert_called_once() + assert env.segment_results == list(result.segments) + env._end_demo_episode_recording.assert_called_once_with(result=result) + + +def test_explicit_candidate_segments_and_actions_remain_lazy() -> None: + env = _SegmentedEnv() + env.create_demo_segments = Mock(side_effect=AssertionError("must not replan")) + + def first_actions(): + assert env.actions == [] + yield 1 + assert env.actions == [1] + yield 2 + + def candidates(): + yield DemoSegment(actions=first_actions(), name="first") + assert env.actions == [1, 2] + yield DemoSegment(actions=(3,), name="second") + + result = execute_demo_episode(env, segments=candidates()) + + assert result.all_success + assert env.actions == [1, 2, 3] + env.create_demo_segments.assert_not_called() + + +def test_explicit_candidate_rejects_planning_arguments_before_lifecycle() -> None: + env = _SegmentedEnv() + env.create_demo_segments = Mock() + env._begin_demo_episode_recording = Mock() + + with pytest.raises(ValueError, match="cannot be combined with plan_kwargs"): + execute_demo_episode(env, segments=DemoSegment(actions=(3,)), goal="other") + + env.create_demo_segments.assert_not_called() + env._begin_demo_episode_recording.assert_not_called() + assert env.actions == [] + + +@pytest.mark.parametrize("segments", [42, "candidate", [object()]]) +def test_explicit_candidate_rejects_invalid_segment_source(segments: Any) -> None: + env = _SegmentedEnv() + env.create_demo_segments = Mock() + + with pytest.raises(TypeError, match="segments must.*DemoSegment"): + execute_demo_episode(env, segments=segments) + + env.create_demo_segments.assert_not_called() + assert env.actions == [] + assert not env._demo_no_auto_reset + + +def test_explicit_candidate_validates_later_segments_when_requested() -> None: + env = _SegmentedEnv() + + def candidates(): + yield DemoSegment(actions=(1,)) + assert env.actions == [1] + yield "invalid segment" + + with pytest.raises(TypeError, match="segments must yield DemoSegment"): + execute_demo_episode(env, segments=candidates()) + + assert env.actions == [1] + assert not env._demo_no_auto_reset + + +def test_empty_explicit_candidates_do_not_fall_back_to_task_planning() -> None: + env = _SegmentedEnv() + env.create_demo_segments = Mock() + + result = execute_demo_episode(env, segments=()) + + env.create_demo_segments.assert_not_called() + assert env.actions == [] + assert result.length == 0 + assert not result.all_success + + +def test_none_segments_preserves_forwarding_legacy_planning_arguments() -> None: + env = _SegmentedEnv() + env.create_demo_segments = Mock(return_value=DemoSegment(actions=(3,))) + + result = execute_demo_episode(env, segments=None, target_uid="object_a") + + env.create_demo_segments.assert_called_once_with(target_uid="object_a") + assert env.actions == [3] + assert result.all_success + + class _LifecycleMetadataEnv: """Populate one shared metadata mapping at lazy lifecycle boundaries.""" @@ -472,11 +591,15 @@ def step(self, action: int): ) -def test_action_generator_failure_safe_stops_before_propagating() -> None: +@pytest.mark.parametrize("explicit_segments", [False, True]) +def test_action_generator_failure_safe_stops_before_propagating( + explicit_segments: bool, +) -> None: env = _GeneratorFailureEnv() + segments = env.create_demo_segments() if explicit_segments else None with pytest.raises(RuntimeError, match="action generation") as error: - execute_demo_episode(env) + execute_demo_episode(env, segments=segments) assert isinstance(error.value.__cause__, ValueError) assert env.actions == [1, 0] @@ -551,11 +674,15 @@ def is_task_success(self) -> torch.Tensor: return torch.ones(self.num_envs, dtype=torch.bool) -def test_execute_demo_episode_supports_staggered_vector_success() -> None: +@pytest.mark.parametrize("explicit_segments", [False, True]) +def test_execute_demo_episode_supports_staggered_vector_success( + explicit_segments: bool, +) -> None: """A completed row freezes while unfinished rows continue their shared plan.""" env = _StaggeredVectorEnv() - result = execute_demo_episode(env) + segments = env.create_demo_segments() if explicit_segments else None + result = execute_demo_episode(env, segments=segments) assert env.actions == [1, 2, 3] assert env.masked_actions == [(2, (False, True)), (3, (False, True))] @@ -849,11 +976,17 @@ def _validate(self) -> torch.Tensor: return torch.ones(1, dtype=torch.bool) -def test_cancellation_after_last_action_does_not_advance_lazy_plan() -> None: +@pytest.mark.parametrize("explicit_segments", [False, True]) +def test_cancellation_after_last_action_does_not_advance_lazy_plan( + explicit_segments: bool, +) -> None: """Cancellation is observed before validation or requesting another segment.""" env = _CancellationEnv() - result = execute_demo_episode(env, should_stop=lambda: bool(env.actions)) + segments = env.create_demo_segments() if explicit_segments else None + result = execute_demo_episode( + env, segments=segments, should_stop=lambda: bool(env.actions) + ) assert env.actions == [1] assert not env.validator_called diff --git a/tests/gym/envs/test_fixed_scene_preparation.py b/tests/gym/envs/test_fixed_scene_preparation.py new file mode 100644 index 000000000..b13983be8 --- /dev/null +++ b/tests/gym/envs/test_fixed_scene_preparation.py @@ -0,0 +1,523 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Controlled Gym preparation without a simulator or implicit reset events.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from tensordict import TensorDict + +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnv +from embodichain.lab.gym.envs.managers.record import record_camera_data +from embodichain.lab.gym.utils.profiler import EnvProfiler +from embodichain.lab.sim.motion.trajectory_augmentation import ( + SceneCase, + ValidationCheck, + ValidationResult, +) +from embodichain.lab.trajectory_generation.initial_state import ( + FixedSceneHost, + InitialStateProfile, +) + + +class _PreparationEnv(EmbodiedEnv): + """Exercise real preparation, reset, and step with CPU host ports.""" + + def __init__(self) -> None: + self.calls: list[str] = [] + self._num_envs = 2 + self.cfg = SimpleNamespace( + events=True, + observations=True, + rewards=True, + dataset=True, + trajectory_auto_save=False, + ignore_terminations=False, + sim_steps_per_control=1, + seed=None, + ) + self.sim_cfg = SimpleNamespace(physics_dt=0.01) + self.sim = SimpleNamespace( + device=torch.device("cpu"), + num_envs=2, + update=Mock(), + reset_objects_state=Mock(), + capture_visualization_safely=Mock(), + ) + self.robot = object() + self._profiler = EnvProfiler(None, torch.device("cpu")) + self._elapsed_steps = torch.full((2,), 7, dtype=torch.long) + self._task_success = torch.ones(2, dtype=torch.bool) + self.episode_success_status = torch.ones(2, dtype=torch.bool) + self._detached_uids_for_reset = [] + self.max_episode_steps = 1 + self.physical_value = torch.tensor([[90.0], [91.0]]) + self.camera = object.__new__(record_camera_data) + self.camera.discard_and_clear = Mock( + side_effect=lambda **kwargs: self.calls.append("camera_discard") + ) + self.camera.save_and_clear = Mock() + self.camera.finalize = Mock() + self.event_manager = SimpleNamespace( + _mode_functor_cfgs={"interval": [SimpleNamespace(func=self.camera)]}, + available_modes=["reset"], + active_functors={}, + apply=Mock(side_effect=lambda **kwargs: self.calls.append("event_reset")), + set_seed=Mock(), + reset=Mock(), + ) + self.observation_manager = SimpleNamespace( + reset=Mock(side_effect=lambda **kwargs: self.calls.append("obs_reset")) + ) + self.reward_manager = SimpleNamespace( + reset=Mock(side_effect=lambda **kwargs: self.calls.append("reward_reset")) + ) + self.dataset_manager = SimpleNamespace( + reset=Mock(side_effect=lambda **kwargs: self.calls.append("dataset_reset")), + apply=Mock(), + available_modes=["save"], + save_failed_episodes=True, + ) + self.rollout_buffer = TensorDict( + { + "obs": TensorDict({"state": torch.full((2, 3, 1), -1.0)}, [2, 3]), + "valid": torch.ones(2, 3, dtype=torch.bool), + }, + [2, 3], + device="cpu", + ) + self._rollout_buffer_mode = "expert" + self._max_rollout_steps = 3 + self.rollout_steps = torch.full((2,), 2, dtype=torch.long) + self.current_rollout_step = 2 + self._traj_steps = torch.full((2,), 2, dtype=torch.long) + self._traj_buffer = None + self._traj_raw_action = object() + self._demo_active_segment_ids = torch.full((2,), 3, dtype=torch.long) + self._demo_active_mask = torch.zeros(2, dtype=torch.bool) + self._demo_segment_participants = torch.ones(2, dtype=torch.bool) + self._demo_active_segment_start_steps = torch.full((2,), 2, dtype=torch.long) + self._demo_active_rollout_start_steps = torch.full((2,), 2, dtype=torch.long) + self._demo_steps = torch.full((2,), 2, dtype=torch.long) + self._demo_episode_metadata = [{"completed": True}, {"completed": True}] + self._active_task_program_bridge = object() + self._demo_no_auto_reset = False + self._replay_no_auto_reset = False + self._np_random = np.random.default_rng(17) + + def get_obs(self, **kwargs): + self.calls.append("get_obs") + return TensorDict({"state": self.physical_value.clone()}, [2], device="cpu") + + def compute_task_state(self, **kwargs): + self.calls.append("task_state") + return torch.zeros(2, dtype=torch.bool), torch.zeros(2, dtype=torch.bool), {} + + def _seed_recording_state(self, obs, env_ids): + self.calls.append("seed_recording") + super()._seed_recording_state(obs, env_ids) + + def _preprocess_action(self, action): + return action + + def _step_action(self, action): + return action + + def _postprocess_action(self, action): + return action + + def get_reward(self, **kwargs): + return torch.zeros(2) + + def _extend_reward(self, rewards, **kwargs): + return rewards + + def _hook_after_sim_step(self, **kwargs): + pass + + def check_truncated(self, **kwargs): + return torch.zeros(2, dtype=torch.bool) + + def is_task_success(self, **kwargs): + return torch.zeros(2, dtype=torch.bool) + + +def _prepare(env: _PreparationEnv, owner: object, *, verification=None): + def prepare(): + env.calls.append("prepare") + assert env._active_task_program_bridge is None + assert env._elapsed_steps.tolist() == [0, 0] + assert not env.rollout_buffer["valid"].any() + + def restore(): + env.calls.append("restore") + env.physical_value = torch.tensor([[1.0], [2.0]]) + + def settle(): + env.calls.append("settle") + env.physical_value += 1.0 + + def verify(): + env.calls.append("verify") + if verification is not None: + return verification + return ValidationResult((ValidationCheck("initial_state", "passed"),)) + + return env.prepare_generation_episode( + owner, prepare=prepare, restore=restore, settle=settle, verify=verify + ) + + +def test_preparation_orders_cleanup_restore_verification_and_first_frame() -> None: + env = _PreparationEnv() + owner = object() + env.acquire_generation_lease(owner) + rng_state = env.np_random.bit_generator.state + torch_rng_state = torch.random.get_rng_state().clone() + + obs, info = _prepare(env, owner) + + assert env.calls == [ + "camera_discard", + "prepare", + "restore", + "settle", + "obs_reset", + "reward_reset", + "dataset_reset", + "verify", + "get_obs", + "task_state", + "seed_recording", + ] + assert obs["state"].tolist() == [[2.0], [3.0]] + assert torch.equal(env.rollout_buffer["obs", "state"][:, 0], obs["state"]) + assert info["elapsed_steps"].tolist() == [0, 0] + for counter in (env.rollout_steps, env._traj_steps, env._demo_steps): + assert counter.tolist() == [0, 0] + assert env.current_rollout_step == 0 + assert not env.episode_success_status.any() + assert not env._task_success.any() + assert not env.rollout_buffer["valid"].any() + assert env._demo_active_mask.all() + assert not env._demo_segment_participants.any() + assert all(not metadata["completed"] for metadata in env._demo_episode_metadata) + assert env._traj_raw_action is None + assert env.np_random.bit_generator.state == rng_state + assert torch.equal(torch.random.get_rng_state(), torch_rng_state) + env.event_manager.apply.assert_not_called() + env.event_manager.reset.assert_not_called() + env.event_manager.set_seed.assert_not_called() + env.dataset_manager.apply.assert_not_called() + env.camera.save_and_clear.assert_not_called() + env.camera.finalize.assert_not_called() + env.sim.update.assert_not_called() + env.sim.reset_objects_state.assert_not_called() + assert env.observation_manager.reset.call_args.kwargs["env_ids"].tolist() == [0, 1] + + +def test_lease_identity_blocks_foreign_calls_without_side_effects() -> None: + env = _PreparationEnv() + owner = object() + env.acquire_generation_lease(owner) + epoch = env.generation_epoch + env.acquire_generation_lease(owner) + assert env.generation_epoch == epoch + assert env.sim._trajectory_generation_owner is owner + for operation in ( + lambda: env.acquire_generation_lease(object()), + lambda: env.release_generation_lease(object()), + lambda: _prepare(env, object()), + lambda: env.reset(seed=9), + lambda: env.step(torch.zeros(2, 1)), + ): + with pytest.raises(RuntimeError): + operation() + assert env.generation_epoch == epoch + assert env.sim._trajectory_generation_owner is owner + assert env.calls == [] + env.event_manager.set_seed.assert_not_called() + env.sim.update.assert_not_called() + + +@pytest.mark.parametrize( + "status", ["failed", "unavailable", "not_run", "empty", "wrong"] +) +def test_unverified_preparation_does_not_publish_first_frame_or_allow_step( + status, +) -> None: + env = _PreparationEnv() + owner = object() + env.acquire_generation_lease(owner) + previous_epoch = env.generation_epoch + validation = ( + object() + if status == "wrong" + else ( + ValidationResult(()) + if status == "empty" + else ValidationResult((ValidationCheck("initial_state", status),)) + ) + ) + with pytest.raises((RuntimeError, TypeError), match="verification"): + _prepare(env, owner, verification=validation) + assert env.generation_epoch > previous_epoch + assert "get_obs" not in env.calls + assert "seed_recording" not in env.calls + assert not env._generation_preparing + assert not env.rollout_buffer["valid"].any() + with pytest.raises(RuntimeError, match="before stepping"): + env.step(torch.zeros(2, 1)) + _prepare(env, owner) + assert env._generation_prepared + + +def test_failed_restore_invalidates_previous_prepared_epoch() -> None: + env = _PreparationEnv() + owner = object() + env.acquire_generation_lease(owner) + _prepare(env, owner) + old_epoch = env.generation_epoch + env.calls.clear() + restore = Mock(side_effect=ValueError("restore failed")) + settle, verify = Mock(), Mock() + with pytest.raises(ValueError, match="restore failed"): + env.prepare_generation_episode( + owner, prepare=Mock(), restore=restore, settle=settle, verify=verify + ) + assert env.generation_epoch > old_epoch + assert not env._generation_prepared + assert "seed_recording" not in env.calls + settle.assert_not_called() + verify.assert_not_called() + with pytest.raises(RuntimeError, match="before stepping"): + env.step(torch.zeros(2, 1)) + + +def test_generation_disables_auto_reset_until_release_without_touching_demo_flags() -> ( + None +): + env = _PreparationEnv() + owner = object() + env.acquire_generation_lease(owner) + _prepare(env, owner) + env.step(torch.zeros(2, 1)) + assert env._elapsed_steps.tolist() == [1, 1] + env.sim.reset_objects_state.assert_not_called() + old_epoch = env.generation_epoch + env.calls.clear() + env._demo_no_auto_reset = True + env.release_generation_lease(owner) + assert env.generation_epoch > old_epoch + assert env._demo_no_auto_reset + assert not env._generation_no_auto_reset + assert env.calls == [] + env.sim.reset_objects_state.assert_not_called() + env._demo_no_auto_reset = False + env.step(torch.zeros(2, 1)) + env.sim.reset_objects_state.assert_called_once() + assert env._elapsed_steps.tolist() == [0, 0] + + +def test_generation_lease_rejects_none_owner() -> None: + env = _PreparationEnv() + with pytest.raises(ValueError, match="None"): + env.acquire_generation_lease(None) + with pytest.raises(RuntimeError, match="lease"): + env.release_generation_lease(None) + + +class _PhysicalAdapter: + """Fake physical port connected to the real Gym and host lifecycle.""" + + def __init__(self, env: _PreparationEnv) -> None: + self.env = env + self.sim = env.sim + self.robot = env.robot + self.accept_state = True + + def signature(self) -> str: + return "fixed-two-row-robot" + + def capture(self) -> torch.Tensor: + return self.env.physical_value.clone() + + def restore(self, state: torch.Tensor) -> None: + self.env.physical_value.copy_(state) + + def verify(self, state: torch.Tensor) -> ValidationResult: + accepted = self.accept_state and torch.equal(self.env.physical_value, state) + return ValidationResult( + (ValidationCheck("physical_initial", "passed" if accepted else "failed"),) + ) + + +def _host_fixture(env: _PreparationEnv): + adapter = _PhysicalAdapter(env) + profile = InitialStateProfile( + profile_id="fixed-test", + prepare=lambda: None, + signature=lambda: "fixed-conditions", + verify=lambda cases: ValidationResult( + (ValidationCheck("task_initial", "passed"),) + ), + settling_steps=1, + ) + cases = tuple( + SceneCase(f"scene-{row}", f"initial-{row}", f"signature-{row}", "task", "robot") + for row in range(2) + ) + return adapter, profile, cases + + +def test_host_and_gym_share_settled_first_frame_and_current_epoch() -> None: + env = _PreparationEnv() + adapter, profile, cases = _host_fixture(env) + env.sim.update.side_effect = lambda dt, steps: env.physical_value.add_(0.5) + with FixedSceneHost(adapter, profile, env=env) as host: + first = host.acquire_case(cases) + settled = env.physical_value.clone() + assert settled.tolist() == [[90.5], [91.5]] + assert torch.equal(env.rollout_buffer["obs", "state"][:, 0], settled) + assert first.epoch == env.generation_epoch + assert host.verify_initial(first).accepted + calls_before_observation = list(env.calls) + observation = host.initial_observation(first) + assert torch.equal(observation["state"], settled) + observation["state"].zero_() + assert torch.equal(host.initial_observation(first)["state"], settled) + assert env.calls == calls_before_observation + env.sim.update.side_effect = None + env.physical_value += 10.0 + restored = host.restore_initial() + assert torch.equal(env.physical_value, settled) + assert torch.equal(env.rollout_buffer["obs", "state"][:, 0], settled) + assert restored.epoch == env.generation_epoch > first.epoch + with pytest.raises(RuntimeError, match="obsolete"): + host.assert_current(first) + with pytest.raises(RuntimeError, match="obsolete"): + host.initial_observation(first) + host.assert_current(restored) + env.sim.reset_objects_state.assert_not_called() + env.event_manager.apply.assert_not_called() + assert not env._generation_no_auto_reset + + +@pytest.mark.parametrize("operation", ["acquire", "restore"]) +def test_host_physical_verification_failure_never_seeds_gym_episode(operation) -> None: + env = _PreparationEnv() + adapter, profile, cases = _host_fixture(env) + with FixedSceneHost(adapter, profile, env=env) as host: + previous = host.acquire_case(cases) if operation == "restore" else None + adapter.accept_state = False + env.calls.clear() + with pytest.raises(RuntimeError, match="verification"): + if operation == "restore": + host.restore_initial() + else: + host.acquire_case(cases) + assert "get_obs" not in env.calls + assert "seed_recording" not in env.calls + assert not env.rollout_buffer["valid"].any() + with pytest.raises(RuntimeError, match="before stepping"): + env.step(torch.zeros(2, 1)) + if previous is not None: + with pytest.raises(RuntimeError, match="obsolete"): + host.assert_current(previous) + adapter.accept_state = True + current = ( + host.restore_initial() + if operation == "restore" + else host.acquire_case(cases) + ) + host.assert_current(current) + + +def test_host_verification_observes_reset_standard_manager_state() -> None: + env = _PreparationEnv() + adapter, profile, cases = _host_fixture(env) + env.reward_manager.dirty = True + env.reward_manager.reset.side_effect = lambda **kwargs: setattr( + env.reward_manager, "dirty", False + ) + profile.verify = lambda cases: ValidationResult( + ( + ValidationCheck( + "manager_initial", "failed" if env.reward_manager.dirty else "passed" + ), + ) + ) + with FixedSceneHost(adapter, profile, env=env) as host: + initial = host.acquire_case(cases) + assert host.verify_initial(initial).accepted + env.reward_manager.dirty = True + restored = host.restore_initial() + assert host.verify_initial(restored).accepted + assert env.reward_manager.reset.call_count == 2 + + +def test_gym_lease_prevents_a_second_pure_sim_host() -> None: + env = _PreparationEnv() + adapter, profile, _ = _host_fixture(env) + owner = object() + env.acquire_generation_lease(owner) + try: + with pytest.raises(RuntimeError, match="already has"): + FixedSceneHost(adapter, profile) + assert env._generation_lease_owner is owner + assert env.sim._trajectory_generation_owner is owner + finally: + env.release_generation_lease(owner) + assert env.sim._trajectory_generation_owner is None + with FixedSceneHost(adapter, profile): + pass + + +def test_pure_sim_host_prevents_a_second_gym_lease() -> None: + env = _PreparationEnv() + adapter, profile, _ = _host_fixture(env) + with FixedSceneHost(adapter, profile) as host: + with pytest.raises(RuntimeError, match="already has"): + env.acquire_generation_lease(object()) + assert env.sim._trajectory_generation_owner is host + assert getattr(env, "_generation_lease_owner", None) is None + owner = object() + env.acquire_generation_lease(owner) + env.release_generation_lease(owner) + + +def test_foreign_sim_owner_cannot_be_cleared_by_stale_gym_lease() -> None: + env = _PreparationEnv() + owner, foreign = object(), object() + env.acquire_generation_lease(owner) + old_epoch = env.generation_epoch + env.sim._trajectory_generation_owner = foreign + for operation in ( + lambda: env.acquire_generation_lease(owner), + lambda: env.release_generation_lease(owner), + ): + with pytest.raises(RuntimeError, match="does not own"): + operation() + assert env._generation_lease_owner is owner + assert env.sim._trajectory_generation_owner is foreign + assert env.generation_epoch == old_epoch diff --git a/tests/lab/task_program/semantics/test_scene.py b/tests/lab/task_program/semantics/test_scene.py index 2d250f52b..8a3972e63 100644 --- a/tests/lab/task_program/semantics/test_scene.py +++ b/tests/lab/task_program/semantics/test_scene.py @@ -31,7 +31,7 @@ ObservedArticulationJointState, SceneSnapshot, ) -from embodichain.lab.sim.planners.base_planner import CollisionWorldInfo +from embodichain.lab.sim.motion.planners.base_planner import CollisionWorldInfo from embodichain.lab.task_program.semantics import ( AmbiguousSceneAffordanceError, GRASP_AFFORDANCE_CAPABILITY, diff --git a/tests/lab/task_program/semantics/test_scene_curobo_integration.py b/tests/lab/task_program/semantics/test_scene_curobo_integration.py index 1c9d1afd4..e3577c0f8 100644 --- a/tests/lab/task_program/semantics/test_scene_curobo_integration.py +++ b/tests/lab/task_program/semantics/test_scene_curobo_integration.py @@ -20,7 +20,7 @@ import torch -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( CuroboPlanOptions, CuroboPlanner, CuroboPlannerCfg, diff --git a/tests/lab/task_program/test_semantic_executor_curobo_gpu.py b/tests/lab/task_program/test_semantic_executor_curobo_gpu.py index 4f32290ae..a4df56afd 100644 --- a/tests/lab/task_program/test_semantic_executor_curobo_gpu.py +++ b/tests/lab/task_program/test_semantic_executor_curobo_gpu.py @@ -50,8 +50,11 @@ from embodichain.lab.sim.atomic_actions.tracking import TrackingPolicy # noqa: E402 from embodichain.lab.sim.cfg import RigidBodyAttributesCfg # noqa: E402 from embodichain.lab.sim.objects import RigidObjectCfg # noqa: E402 -from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator # noqa: E402 -from embodichain.lab.sim.planners.curobo.curobo_planner import ( # noqa: E402 +from embodichain.lab.sim.motion.planners import ( + MotionGenCfg, + MotionGenerator, +) # noqa: E402 +from embodichain.lab.sim.motion.planners.curobo.curobo_planner import ( # noqa: E402 CuroboAutoGenCfg, CuroboPlannerCfg, CuroboWorldCfg, diff --git a/tests/lab/trajectory_generation/__init__.py b/tests/lab/trajectory_generation/__init__.py new file mode 100644 index 000000000..3818cfd8e --- /dev/null +++ b/tests/lab/trajectory_generation/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__ = [] diff --git a/tests/lab/trajectory_generation/test_atomic_source.py b/tests/lab/trajectory_generation/test_atomic_source.py new file mode 100644 index 000000000..e10ca2904 --- /dev/null +++ b/tests/lab/trajectory_generation/test_atomic_source.py @@ -0,0 +1,103 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Atomic export preserves source ownership, phase boundaries and passive geometry.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.lab.trajectory_generation.integrations.atomic import ( + export_pickup_templates, +) + + +def _inputs(): + q = torch.zeros(2, 9, 3) + q[:, :3, 0] = torch.tensor([0.0, 0.1, 0.2]) + q[:, 3:, 0] = 0.2 + q[:, 5:, 1] = 0.03 + dt = torch.full((2, 9), 0.05, dtype=torch.float64) + dt[:, (0, 3)] = 0.0 + compiled = SimpleNamespace( + action_plans=( + SimpleNamespace(skill_id="move_end_effector"), + SimpleNamespace(skill_id="pick_up"), + ), + plan_success=torch.ones(2, dtype=torch.bool), + trajectory=SimpleNamespace(positions=q, dt=dt, env_ids=torch.arange(2)), + action_waypoint_offset=lambda index: 3, + segment=lambda index, name: { + "approach": SimpleNamespace(start=3, stop=5), + "close": SimpleNamespace(start=5, stop=7), + "lift": SimpleNamespace(start=7, stop=9), + }[name], + ) + robot = SimpleNamespace( + num_instances=2, + dof=3, + joint_names=("arm", "finger", "mimic"), + mimic_ids=(2,), + mimic_parents=(1,), + mimic_multipliers=(-1.0,), + mimic_offsets=(0.0,), + get_joint_ids=lambda name: (0,), + ) + return compiled, robot + + +def test_export_keeps_atomic_commands_and_declares_real_hold_and_mimic_geometry(): + compiled, robot = _inputs() + originals = compiled.trajectory.positions.clone(), compiled.trajectory.dt.clone() + templates = export_pickup_templates(compiled, robot, control_dt=0.05, hold_steps=3) + assert len(templates) == 2 + for template in templates: + assert template.positions.shape == (12, 3) + torch.testing.assert_close(template.positions[:9, :2], originals[0][0, :, :2]) + torch.testing.assert_close(template.positions[:, 2], -template.positions[:, 1]) + assert template.dt[0] == 0 and torch.all(template.dt[1:] == 0.05) + assert template.phases[0].allowed_operators == ("joint_residual",) + assert all(not p.allowed_operators for p in template.phases[1:]) + assert template.phases[-1].start_index == 9 + assert torch.equal(compiled.trajectory.positions, originals[0]) + assert torch.equal(compiled.trajectory.dt, originals[1]) + templates[0].positions.zero_() + assert templates[1].positions.abs().sum() > 0 + + +@pytest.mark.parametrize( + "failure", + ["failed_plan", "boundary", "clock", "row_order", "skill", "missing_segment"], +) +def test_export_rejects_incompatible_atomic_plan(failure): + compiled, robot = _inputs() + if failure == "failed_plan": + compiled.plan_success[0] = False + if failure == "boundary": + compiled.trajectory.positions[:, 3, 0] += 0.1 + if failure == "clock": + compiled.trajectory.dt[:, 5] = 0.1 + if failure == "row_order": + compiled.trajectory.env_ids = torch.tensor([1, 0]) + if failure == "skill": + compiled.action_plans[1].skill_id = "place" + if failure == "missing_segment": + compiled.segment = lambda index, name: SimpleNamespace(start=0, stop=1) + with pytest.raises(ValueError): + export_pickup_templates(compiled, robot, control_dt=0.05) diff --git a/tests/lab/trajectory_generation/test_contact.py b/tests/lab/trajectory_generation/test_contact.py new file mode 100644 index 000000000..c3fc669c8 --- /dev/null +++ b/tests/lab/trajectory_generation/test_contact.py @@ -0,0 +1,268 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Physical contact evidence, phase permissions and complete collision geometry.""" + +from __future__ import annotations + +from dataclasses import replace +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CandidateIdentity, + CandidateTrajectoryBatch, + MotionSnapshot, + SceneCase, + TrajectoryPhase, +) +from embodichain.lab.trajectory_generation.integrations.contact import ( + PickUpContactProfile, + PickUpMotionValidator, +) + + +def _monitor(): + validator = object.__new__(PickUpMotionValidator) + validator.profile = PickUpContactProfile() + validator.world = SimpleNamespace(adjacent_pairs=set()) + validator.robot = SimpleNamespace(num_instances=2) + validator._users = { + 1: (0, "cube"), + 2: (0, "gripper_finger1_link_1"), + 3: (0, "gripper_finger2_link_1"), + 4: (0, "bench"), + 5: (0, "forearm"), + 11: (1, "cube"), + 12: (1, "gripper_finger1_link_1"), + } + data = np.zeros((2, 11)) + data[:, 9] = 0.001 + native = {"data": data, "users": np.array([[1, 2], [3, 1]])} + validator._physics = SimpleNamespace( + get_cpu_contact_buffer=lambda: (native["data"], native["users"]) + ) + pose = torch.eye(4).repeat(2, 1, 1) + pose[:, 2, 3] = 0.505 + observed = {"object_pose": pose.clone(), "tcp_pose": pose.clone()} + validator.observations = lambda: observed + phases = ( + TrajectoryPhase("transit", 0, 1), + TrajectoryPhase("approach", 1, 2, kind="contact"), + TrajectoryPhase("close", 2, 3, kind="contact"), + TrajectoryPhase("lift", 3, 4, kind="contact"), + TrajectoryPhase("hold", 4, 6, kind="hold"), + ) + cases = [ + SceneCase(f"case_{i}", "initial", "scene", "pickup", "robot") for i in range(2) + ] + initial = torch.eye(4) + initial[2, 3] = 0.325 + snapshots = [ + MotionSnapshot( + case, + ("joint",), + torch.zeros(1), + torch.zeros(1), + torch.eye(4), + {"cube": initial}, + ) + for case in cases + ] + candidates = [ + CandidateTrajectoryBatch( + torch.zeros(1, 6, 1), + torch.tensor([[0.0, 0.05, 0.05, 0.05, 0.05, 0.05]]), + torch.tensor([6]), + ( + CandidateIdentity( + case.scene_case_id, + "initial", + f"candidate_{i}", + f"family_{i}", + "source", + "v1", + "template", + ), + ), + ("joint",), + (phases,), + source_row_indices=torch.tensor([i]), + ) + for i, case in enumerate(cases) + ] + validator.begin_rollout(candidates, snapshots) + return validator, native, observed + + +def _hold(validator): + for _ in range(5): + validator.observe_substep(4, (True, False), physics_dt=0.25) + return validator.rollout_validation(0) + + +def test_stable_lift_requires_both_real_finger_contacts(): + validator, native, _ = _monitor() + assert _hold(validator).accepted + validator, native, _ = _monitor() + native["data"][1, 9] = 0.0 # Geometric proximity without normal impulse. + checks = {c.check_id: c for c in _hold(validator).checks} + assert checks["physical_contacts"].status == "passed" + assert checks["held_object_stability"].status == "failed" + assert checks["held_object_stability"].metrics["finger_1_contact_fraction"] == 0 + + +@pytest.mark.parametrize( + "failure", ["table", "robot", "unknown", "cross_row", "penetration"] +) +def test_forbidden_contact_rejects_even_a_stable_lift(failure): + validator, native, _ = _monitor() + if failure == "table": + native["users"][0] = (1, 4) + if failure == "robot": + native["users"][0] = (1, 5) + if failure == "unknown": + native["users"][0] = (1, 9999) + if failure == "cross_row": + native["users"][0] = (1, 11) + if failure == "penetration": + native["data"][0, 10] = -0.003 + checks = {c.check_id: c for c in _hold(validator).checks} + assert checks["physical_contacts"].status == "failed" + + +@pytest.mark.parametrize("failure", ["slip", "rotate", "fall", "empty", "short"]) +def test_hold_rejects_slipping_rotating_falling_and_empty_grasps(failure): + validator, _, observed = _monitor() + validator.observe_substep(4, (True, False), physics_dt=0.25) + if failure == "slip": + observed["object_pose"][0, 0, 3] += 0.02 + if failure == "rotate": + a = 0.25 + observed["object_pose"][0, :2, :2] = torch.tensor( + [[np.cos(a), -np.sin(a)], [np.sin(a), np.cos(a)]] + ) + if failure == "fall": + observed["object_pose"][0, 2, 3] = 0.325 + if failure == "empty": + observed["object_pose"][0, 0, 3] += 0.1 + result = validator.rollout_validation(0) if failure == "short" else _hold(validator) + assert result.checks[1].status == "failed" + + +def test_contact_permissions_depend_on_phase_and_grasp_entry_distance(): + validator, _, _ = _monitor() + finger = validator.profile.finger_links[0] + assert not validator._allowed("cube", finger, "transit", 0.0, 0.01) + assert not validator._allowed("cube", finger, "approach", 0.0, 0.05) + assert validator._allowed("cube", finger, "approach", 0.0, 0.01) + assert validator._allowed("cube", finger, "close", 0.0) + assert not validator._allowed("cube", "forearm", "close", 0.0) + assert validator._allowed("cube", "bench", "lift", 0.001) + assert not validator._allowed("cube", "bench", "hold", 0.001) + + +@pytest.mark.parametrize("failure", ["overflow", "nan", "shape"]) +def test_missing_or_unbounded_native_evidence_raises(failure): + validator, native, _ = _monitor() + if failure == "overflow": + validator.profile.max_contacts_per_step = 1 + if failure == "nan": + native["data"][0, 9] = np.nan + if failure == "shape": + native["users"] = np.zeros((1, 2)) + with pytest.raises(ValueError, match="Native contact evidence"): + _hold(validator) + + +def test_unobserved_hold_cannot_pass(): + validator, _, _ = _monitor() + assert not validator.rollout_validation(0).accepted + + +def test_full_state_collision_uses_collision_shapes_and_moving_fingers(tmp_path): + pytest.importorskip("fcl") + pytest.importorskip("yourdfpy") + import pytorch_kinematics as pk + from embodichain.lab.sim.shapes import CubeCfg + from embodichain.lab.trajectory_generation.integrations._collision import ( + _FullStateCollisionWorld, + ) + + # Deliberately no visual mesh: a render-geometry fallback would miss this finger. + urdf = """ + + """ + path = tmp_path / "robot.urdf" + path.write_text(urdf) + obj = SimpleNamespace( + cfg=SimpleNamespace(shape=CubeCfg(size=(0.05, 0.05, 0.05))), + get_body_scale=lambda: torch.ones(1, 3), + ) + sim = SimpleNamespace( + get_rigid_object_uid_list=lambda: ("cube", "bench"), + get_rigid_object=lambda uid: obj, + ) + robot = SimpleNamespace( + cfg=SimpleNamespace(body_scale=(1.0, 1.0, 1.0), fpath=str(path)), + device=torch.device("cpu"), + joint_names=("slide",), + pk_chain=pk.build_chain_from_urdf(urdf), + ) + world = _FullStateCollisionWorld(sim, robot) + root = torch.eye(4) + root[2, 3] = 1.0 + poses = world.link_poses(torch.tensor([[0.0], [0.1]]), root) + cube = np.eye(4) + cube[:3, 3] = (0.2, 0.0, 1.0) + bench = np.eye(4) + bench[:3, 3] = (0.5, 0.0, 1.0) + objects = {"cube": cube, "bench": bench} + assert ( + world.collisions( + {"finger": poses["finger"][0]}, objects, "cube", lambda a, b: False + ) + == "finger / cube" + ) + assert ( + world.collisions( + {"finger": poses["finger"][1]}, objects, "cube", lambda a, b: False + ) + is None + ) + objects["cube"] = bench.copy() # Held object can collide while the arm clears. + assert ( + world.collisions( + {"finger": poses["finger"][1]}, objects, "cube", lambda a, b: False + ) + == "cube / bench" + ) + + +def test_cross_row_contact_rejects_active_second_body_when_first_row_is_idle(): + validator, native, _ = _monitor() + native["users"][0] = (11, 1) + assert _hold(validator).checks[0].status == "failed" + + +def test_slip_during_lift_is_not_hidden_by_a_stable_terminal_hold(): + validator, _, observed = _monitor() + validator.observe_substep(3, (True, False), physics_dt=0.25) + observed["object_pose"][0, 0, 3] += 0.02 + assert _hold(validator).checks[1].status == "failed" diff --git a/tests/lab/trajectory_generation/test_cube_grasp_parallel.py b/tests/lab/trajectory_generation/test_cube_grasp_parallel.py new file mode 100644 index 000000000..cff6fd9e0 --- /dev/null +++ b/tests/lab/trajectory_generation/test_cube_grasp_parallel.py @@ -0,0 +1,132 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Physical cube-pickup augmentation, sustained-grasp checks, and video evidence.""" + +from __future__ import annotations + +import json +from pathlib import Path +import subprocess +import sys + +import numpy as np +import pytest + + +def test_grasp_results_reject_fallen_slipping_and_distant_objects() -> None: + from examples.sim.motion.trajectory_generation.cube_grasp_parallel import ( + _grasp_results, + ) + + cubes = np.tile(np.eye(4), (4, 6, 1, 1)) + cubes[:, 0, 2, 3] = 0.325 + cubes[:, 1:, 2, 3] = 0.505 + tcps = cubes.copy() + cubes[1, -1, 2, 3] = 0.325 # Fell after initially reaching the lift target. + cubes[2, -1, 0, 3] += 0.02 # Slipped relative to the TCP during the hold. + cubes[3, 1:, 0, 3] += 0.1 # Elevated, but never near the gripper. + results = _grasp_results(cubes, tcps, hold_start=3) + assert [result["success"] for result in results] == [True, False, False, False] + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.requires_sim +def test_parallel_cube_pickups_preserve_contact_commands_and_save_synchronized_video( + tmp_path: Path, +) -> None: + av = pytest.importorskip("av") + root = Path(__file__).resolve().parents[3] + output = tmp_path / "parallel" + # Run the real entry point in a child process to exercise native cleanup + # and the exit status as well as the physics and recorder. + process = subprocess.run( + [ + sys.executable, + "examples/sim/motion/trajectory_generation/cube_grasp_parallel.py", + "--output", + str(output), + ], + cwd=root, + capture_output=True, + text=True, + timeout=180, + ) + assert process.returncode == 0, (process.stdout + process.stderr)[-8000:] + report = json.loads((output / "report.json").read_text()) + assert report["success"] + assert [row["yaw_degrees"] for row in report["results"]] == [0, 0, 90, 90] + assert all( + row["success"] and row["min_hold_lift_m"] >= 0.12 for row in report["results"] + ) + assert min(report["max_path_pair_separation_m"]) > 0.03 + with np.load(output / "rollout.npz") as saved: + targets = saved["commanded_qpos"] + reference = saved["reference_qpos"] + transit_stop = report["transit_stop"] + np.testing.assert_array_equal( + targets[:, transit_stop - 1 : reference.shape[1]], + reference[:, transit_stop - 1 :], + ) + assert not np.array_equal( + targets[1, 1 : transit_stop - 1], reference[1, 1 : transit_stop - 1] + ) + assert targets.shape == saved["measured_qpos"].shape + np.testing.assert_allclose( + saved["measured_qpos"][:, 0], + np.repeat(saved["measured_qpos"][0:1, 0], 4, axis=0), + atol=1e-6, + rtol=0, + ) + np.testing.assert_allclose( + saved["cube_poses"][:, 0], + np.repeat(saved["cube_poses"][0:1, 0], 4, axis=0), + atol=1e-5, + rtol=0, + ) + # The quarter-turn changes the closing orientation while preserving + # the centered grasp position and downward approach direction. + grasps = saved["grasp_poses"] + np.testing.assert_allclose( + grasps[:, :3, 3], np.repeat(grasps[0:1, :3, 3], 4, axis=0), atol=1e-6 + ) + assert abs(np.dot(grasps[0, :3, 0], grasps[2, :3, 0])) < 1e-5 + actual_rotations = saved["tcp_poses"][:, -1, :3, :3] + assert abs(np.dot(actual_rotations[0, :, 0], actual_rotations[2, :, 0])) < 0.05 + times = saved["timestamps"] + np.testing.assert_allclose( + times, np.arange(len(times)) / report["fps"], atol=1e-7 + ) + with av.open(str(output / "preview.mp4")) as video: + stream = video.streams.video[0] + assert stream.codec_context.name == "h264" + assert (stream.width, stream.height) == (1280, 1056) + assert stream.average_rate == report["fps"] == 20 + frames = list(video.decode(video=0)) + assert len(frames) == report["frame_count"] == len(times) + np.testing.assert_allclose( + [float(frame.pts * frame.time_base) for frame in frames], times, atol=1e-7 + ) + first = frames[0].to_ndarray(format="rgb24") + last = frames[-1].to_ndarray(format="rgb24") + for row in range(4): + top = (row // 2) * 528 + 48 # Exclude the changing text header. + left = (row % 2) * 640 + initial = first[top : top + 480, left : left + 640].astype(float) + final = last[top : top + 480, left : left + 640].astype(float) + assert initial.std() > 1 + assert np.abs(initial - final).mean() > 1 diff --git a/tests/lab/trajectory_generation/test_episode_sinks.py b/tests/lab/trajectory_generation/test_episode_sinks.py new file mode 100644 index 000000000..a839a5683 --- /dev/null +++ b/tests/lab/trajectory_generation/test_episode_sinks.py @@ -0,0 +1,331 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Real LeRobot episode sealing, causal evidence, and submission retries.""" + +from __future__ import annotations + +import json +from dataclasses import replace +from pathlib import Path + +import numpy as np +import pytest +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CandidateIdentity, + ExpertEpisode, + ValidationCheck, + ValidationResult, +) +from embodichain.lab.trajectory_generation.sinks import LeRobotEpisodeSink + +_FPS = 20 +_STEPS = 3 + + +def _episode(commit: str = "commit-1", *, rgb: bool = False) -> ExpertEpisode: + observations = { + "joint_positions": torch.arange((_STEPS + 1) * 2, dtype=torch.float32).reshape( + _STEPS + 1, 2 + ) + / 10, + "scalar": torch.arange(_STEPS + 1, dtype=torch.int64).reshape(-1, 1), + } + if rgb: + observations["images.front"] = torch.arange( + (_STEPS + 1) * 4 * 5 * 3, dtype=torch.uint8 + ).reshape(_STEPS + 1, 4, 5, 3) + return ExpertEpisode( + identity=CandidateIdentity( + "case", + "initial", + f"candidate-{commit}", + "family", + "source", + "revision", + "template", + ), + observations=observations, + actions=torch.tensor([[0.1], [0.2], [0.3]]), + timestamps=7 + torch.arange(_STEPS + 1, dtype=torch.float64) / _FPS, + action_representation="qpos", + validation=ValidationResult((ValidationCheck("task_success", "passed"),)), + episode_id=f"episode-{commit}", + commit_id=commit, + metadata={"task": "move the cube", "nested": {"actual": True}}, + ) + + +def test_real_lerobot_sealing_preserves_causal_frames_rgb_and_terminal_evidence( + tmp_path: Path, +) -> None: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + episode = _episode(rgb=True) + with LeRobotEpisodeSink(tmp_path / "collection", fps=_FPS) as sink: + receipt = sink.submit(episode) + assert receipt.confirmed, receipt.error + assert sink.drain() == () + shard = Path(receipt.storage_id) + dataset = LeRobotDataset( + repo_id="embodichain/trajectory-generation", root=shard / "dataset" + ) + assert dataset.num_episodes == 1 + assert len(dataset) == _STEPS + assert torch.equal( + dataset[0]["observation.joint_positions"], + episode.observations["joint_positions"][0], + ) + assert dataset[_STEPS - 1]["action"].item() == pytest.approx( + episode.actions[-1].item() + ) + assert dataset[0]["observation.images.front"].shape == (3, 4, 5) + with np.load(shard / "terminal.npz", allow_pickle=False) as evidence: + np.testing.assert_array_equal( + evidence["timestamps"], episode.timestamps.numpy() + ) + np.testing.assert_array_equal( + evidence["observation.joint_positions"], + episode.observations["joint_positions"][-1].numpy(), + ) + np.testing.assert_array_equal( + evidence["observation.images.front"], + episode.observations["images.front"][-1].numpy(), + ) + metadata = json.loads((shard / "episode.json").read_text()) + assert metadata["identity"]["candidate_id"] == episode.identity.candidate_id + assert metadata["metadata"] == {"task": "move the cube", "nested": {"actual": True}} + assert metadata["validation"][0]["status"] == "passed" + manifest = json.loads((shard.parent / "manifest.json").read_text()) + assert [item["commit_id"] for item in manifest["episodes"]] == [episode.commit_id] + + +def test_same_commit_is_idempotent_and_new_commits_continue_after_finalization( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with LeRobotEpisodeSink(tmp_path / "collection", fps=_FPS) as sink: + calls = [] + original = sink._write_dataset + + def write(*args): + calls.append(args[0]) + return original(*args) + + monkeypatch.setattr(sink, "_write_dataset", write) + first = sink.submit(_episode()) + duplicate = sink.submit(_episode(), submission_id=1) + second = sink.submit(_episode("commit-2")) + assert first.confirmed and duplicate.confirmed and second.confirmed + assert first.storage_id == duplicate.storage_id + assert duplicate.submission_id == 1 + assert len(calls) == 2 + manifest = json.loads((sink.root / "manifest.json").read_text()) + assert len(manifest["episodes"]) == 2 + + +@pytest.mark.parametrize( + "failure_stage", + ["_write_dataset", "_write_evidence", "_write_manifest", "_verify_evidence"], +) +def test_failure_receipt_retry_reuses_one_logical_episode( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, failure_stage: str +) -> None: + with LeRobotEpisodeSink(tmp_path / "collection", fps=_FPS) as sink: + original = getattr(sink, failure_stage) + + def fail_after_success(*args): + original(*args) + raise OSError("injected persistence failure") + + monkeypatch.setattr(sink, failure_stage, fail_after_success) + failed = sink.submit(_episode()) + assert not failed.confirmed + assert "injected persistence failure" in failed.error + monkeypatch.setattr(sink, failure_stage, original) + + # The first attempt already sealed the shard even when a later barrier + # failed. Retrying must never create another training episode. + def reject_rewrite(*args): + raise AssertionError("A sealed shard must be reused") + + monkeypatch.setattr(sink, "_write_dataset", reject_rewrite) + accepted = sink.submit(_episode(), submission_id=1) + assert accepted.confirmed, accepted.error + assert accepted.storage_id == failed.storage_id + assert ( + len(json.loads((sink.root / "manifest.json").read_text())["episodes"]) == 1 + ) + + +def test_partial_writer_failure_can_rebuild_same_uncommitted_shard( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + original = LeRobotDataset.add_frame + with LeRobotEpisodeSink(tmp_path / "collection", fps=_FPS) as sink: + + def fail_frame(*args): + raise OSError("frame write failed") + + monkeypatch.setattr(LeRobotDataset, "add_frame", fail_frame) + failed = sink.submit(_episode()) + assert not failed.confirmed + assert not (sink.root / "manifest.json").exists() + monkeypatch.setattr(LeRobotDataset, "add_frame", original) + accepted = sink.submit(_episode(), submission_id=1) + assert accepted.confirmed, accepted.error + assert accepted.storage_id == failed.storage_id + + +def test_changed_payload_under_same_commit_is_rejected(tmp_path: Path) -> None: + with LeRobotEpisodeSink(tmp_path / "collection", fps=_FPS) as sink: + episode = _episode() + assert sink.submit(episode).confirmed + changed = replace(episode, actions=episode.actions + 0.1) + with pytest.raises(ValueError, match="changed episode payload"): + sink.submit(changed, submission_id=1) + assert ( + len(json.loads((sink.root / "manifest.json").read_text())["episodes"]) == 1 + ) + + +@pytest.mark.parametrize( + "observation", + [ + torch.zeros(_STEPS + 1, 0, 2), + torch.zeros(_STEPS + 1, 2, dtype=torch.bool), + torch.zeros(_STEPS + 1, 4, 5, 1, dtype=torch.uint8), + torch.zeros(_STEPS + 1, 4, 5, 3), + ], +) +def test_unsupported_observations_fail_before_shard_write( + tmp_path: Path, observation: torch.Tensor +) -> None: + with LeRobotEpisodeSink(tmp_path / "collection", fps=_FPS) as sink: + episode = replace(_episode(), observations={"unsupported": observation}) + with pytest.raises(ValueError, match="Unsupported observation"): + sink.submit(episode) + assert tuple(sink.root.iterdir()) == (sink.root / ".writer.lock",) + + +def test_nonuniform_clock_size_and_failed_validation_rejected_before_writes( + tmp_path: Path, +) -> None: + with LeRobotEpisodeSink(tmp_path / "collection", fps=_FPS) as sink: + episode = _episode() + times = episode.timestamps.clone() + times[-1] += 0.01 + with pytest.raises(ValueError, match="fixed control clock"): + sink.submit(replace(episode, timestamps=times)) + failed = ValidationResult((ValidationCheck("task_success", "failed"),)) + with pytest.raises(ValueError, match="accepted episodes"): + sink.submit(replace(episode, validation=failed)) + with LeRobotEpisodeSink( + tmp_path / "bounded", fps=_FPS, max_episode_bytes=1 + ) as sink: + with pytest.raises(ValueError, match="max_episode_bytes"): + sink.submit(_episode()) + + +def test_terminal_corruption_cannot_return_confirmed_receipt( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with LeRobotEpisodeSink(tmp_path / "collection", fps=_FPS) as sink: + original = sink._write_evidence + + def corrupt(path, *args): + original(path, *args) + (path / "terminal.npz").write_bytes(b"invalid") + + monkeypatch.setattr(sink, "_write_evidence", corrupt) + receipt = sink.submit(_episode()) + assert not receipt.confirmed + assert not (sink.root / "manifest.json").exists() + + +def test_sink_rejects_existing_collection_and_closed_submissions( + tmp_path: Path, +) -> None: + root = tmp_path / "collection" + sink = LeRobotEpisodeSink(root, fps=_FPS) + with pytest.raises(ValueError, match="new or empty"): + LeRobotEpisodeSink(root, fps=_FPS) + sink.close() + sink.close() + with pytest.raises(RuntimeError, match="closed"): + sink.submit(_episode()) + + +def test_float64_storage_preserves_precision_beyond_torch_reader_defaults( + tmp_path: Path, +) -> None: + precise = torch.tensor( + [[0.123456789123456, 0.987654321987654]], dtype=torch.float64 + ) + episode = replace( + _episode(), + observations={"precise": precise.repeat(_STEPS + 1, 1)}, + actions=precise.repeat(_STEPS, 1), + ) + with LeRobotEpisodeSink(tmp_path / "collection", fps=_FPS) as sink: + receipt = sink.submit(episode) + assert receipt.confirmed, receipt.error + + +def test_confirmed_duplicate_only_reads_existing_artifacts( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with LeRobotEpisodeSink(tmp_path / "collection", fps=_FPS) as sink: + first = sink.submit(_episode()) + assert first.confirmed, first.error + + def no_rewrite(*args): + raise AssertionError("Confirmed duplicates must not write again") + + monkeypatch.setattr(sink, "_write_dataset", no_rewrite) + monkeypatch.setattr(sink, "_write_evidence", no_rewrite) + monkeypatch.setattr(sink, "_write_manifest", no_rewrite) + duplicate = sink.submit(_episode(), submission_id=1) + assert duplicate.confirmed, duplicate.error + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_numeric_pose_matrices_roundtrip_with_original_shape( + tmp_path: Path, dtype: torch.dtype +) -> None: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + poses = torch.eye(4, dtype=dtype).repeat(_STEPS + 1, 1, 1) + poses[:, :3, 3] = torch.arange((_STEPS + 1) * 3, dtype=dtype).reshape(-1, 3) / 7 + episode = replace(_episode(), observations={"object_pose": poses}) + with LeRobotEpisodeSink(tmp_path / "matrix", fps=_FPS) as sink: + receipt = sink.submit(episode) + assert receipt.confirmed, receipt.error + shard = Path(receipt.storage_id) + metadata = json.loads((shard / "episode.json").read_text()) + assert metadata["observation_shapes"]["object_pose"] == [4, 4] + assert metadata["features"]["observation.object_pose"]["shape"] == [16] + dataset = LeRobotDataset( + repo_id="embodichain/trajectory-generation", root=shard / "dataset" + ) + assert dataset[0]["observation.object_pose"].shape == (16,) + with np.load(shard / "terminal.npz", allow_pickle=False) as terminal: + np.testing.assert_array_equal( + terminal["observation.object_pose"], poses[-1].numpy() + ) diff --git a/tests/lab/trajectory_generation/test_execution.py b/tests/lab/trajectory_generation/test_execution.py new file mode 100644 index 000000000..db84bd8fc --- /dev/null +++ b/tests/lab/trajectory_generation/test_execution.py @@ -0,0 +1,684 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Causal execution through real Gym control ports and deterministic CPU physics.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from tensordict import TensorDict + +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnv +from embodichain.lab.gym.utils.profiler import EnvProfiler +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CandidateIdentity, + CandidateTrajectoryBatch, + SceneCase, + TrajectoryPhase, + ValidationCheck, + ValidationResult, +) +from embodichain.lab.trajectory_generation.execution import QposRolloutExecutor + + +def _passed(*args) -> ValidationResult: + return ValidationResult((ValidationCheck("task_success", "passed"),)) + + +class _Robot: + def __init__(self) -> None: + self.joint_names = ("arm", "locked") + self.active_joint_ids = [0] + self.qpos = torch.tensor([[0.0, 0.3], [0.0, 0.3]]) + self.target = self.qpos.clone() + self.target_velocity = torch.zeros_like(self.qpos) + self.force = torch.zeros_like(self.qpos) + self.pose = torch.eye(4).repeat(2, 1, 1) + self.commands = [] + self.offset = 0.0 + + def get_qpos(self, target=False): + return self.target if target else self.qpos + + def get_qpos_limits(self): + return torch.tensor([-1.0, 1.0]).repeat(2, 2, 1) + + def get_qvel(self, target=False): + assert target + return self.target_velocity + + def get_qf(self): + return self.force + + def get_local_pose(self, to_matrix=False): + assert to_matrix + return self.pose + + def set_qpos(self, qpos, *, joint_ids, target=True): + assert target + self.target[:, joint_ids] = qpos + self.offset + self.commands.append(self.target.clone()) + + def set_qvel(self, qvel, *, joint_ids, target=True): + self.target_velocity[:, joint_ids] = qvel + + def set_qf(self, qf, *, joint_ids): + self.force[:, joint_ids] = qf + + +class _Sim: + num_envs = 2 + device = torch.device("cpu") + + def __init__(self, robot): + self.robot = robot + self.simulation_time = 11.0 + self.updates = [] + self.fail_at = None + self.clock_scale = 1.0 + self.after_update = lambda: None + self.cube = SimpleNamespace(pose=torch.eye(4).repeat(2, 1, 1)) + self.cube.get_local_pose = lambda to_matrix: self.cube.pose + + def update(self, dt, steps): + self.updates.append((dt, steps)) + if len(self.updates) == self.fail_at: + raise RuntimeError("physics failed") + self.robot.qpos += 0.5 * (self.robot.target - self.robot.qpos) + self.simulation_time += dt * steps * self.clock_scale + self.after_update() + + def get_rigid_object_uid_list(self): + return ["cube"] + + def get_rigid_object(self, uid): + assert uid == "cube" + return self.cube + + +class _Env(EmbodiedEnv): + """Use BaseEnv.step and normal EmbodiedEnv action/demo machinery unchanged.""" + + def __init__(self, sim, robot): + self.sim, self.robot = sim, robot + self._num_envs = 2 + self.active_joint_ids = [0] + self.cfg = SimpleNamespace( + sim_steps_per_control=2, + events=False, + dataset=False, + ignore_terminations=False, + trajectory_auto_save=False, + ) + self.sim_cfg = SimpleNamespace(physics_dt=0.01) + self.max_episode_steps = 100 + self._profiler = EnvProfiler(None, sim.device) + self._elapsed_steps = torch.zeros(2, dtype=torch.long) + self._traj_buffer = None + self.rollout_buffer = None + self.rollout_steps = torch.zeros(2, dtype=torch.long) + self._demo_steps = torch.zeros(2, dtype=torch.long) + self._demo_active_segment_ids = torch.zeros(2, dtype=torch.long) + self._demo_active_mask = torch.ones(2, dtype=torch.bool) + self._demo_segment_participants = torch.zeros(2, dtype=torch.bool) + self.obs_calls = 0 + self.fail_observation_at = None + self.terminate_after = [None, None] + self.post_target_offset = 0.0 + self.action_manager = Mock(process_action=Mock(side_effect=self._post)) + self.recorded_masks = [] + + def _post(self, action, mode): + assert mode == "post", "ControllerAction must bypass policy preprocessing" + self.robot.target[:, 0] += self.post_target_offset + return action + 8.0 # Postprocessor output is not the command label. + + def get_obs(self, **kwargs): + self.obs_calls += 1 + if self.obs_calls == self.fail_observation_at: + raise RuntimeError("camera unavailable") + return TensorDict({"state": self.robot.qpos.clone()}, batch_size=[2]) + + def get_info(self, **kwargs): + return { + "success": torch.tensor( + [ + limit is not None and len(self.sim.updates) >= limit + for limit in self.terminate_after + ] + ), + "fail": torch.zeros(2, dtype=torch.bool), + } + + def get_reward(self, **kwargs): + return torch.zeros(2) + + def _extend_reward(self, rewards, **kwargs): + return rewards + + def check_truncated(self, **kwargs): + return torch.zeros(2, dtype=torch.bool) + + def is_task_success(self, **kwargs): + return torch.ones(2, dtype=torch.bool) + + def _hook_after_sim_step(self, **kwargs): + self.recorded_masks.append(self._demo_active_mask.clone()) + self._demo_steps += self._demo_active_mask + + +@pytest.fixture(params=[False, True], ids=["pure_sim", "gym"]) +def scene(request): + robot = _Robot() + sim = _Sim(robot) + env = _Env(sim, robot) if request.param else None + cases = tuple( + SceneCase(f"case-{i}", f"initial-{i}", "scene", "task", "robot") + for i in range(2) + ) + binding = SimpleNamespace(cases=cases) + initial = env.get_obs() if env is not None else None + snapshots = tuple( + SimpleNamespace( + root_pose=robot.pose[i].clone(), + entity_poses={"cube": sim.cube.pose[i].clone()}, + ) + for i in range(2) + ) + host = SimpleNamespace( + adapter=SimpleNamespace(robot=robot, sim=sim, atol=1e-5), + env=env, + profile=SimpleNamespace(physics_dt=0.01), + assert_current=Mock(), + verify_initial=Mock(return_value=_passed()), + snapshots=Mock(return_value=snapshots), + initial_observation=lambda value: ( + initial.clone() if initial is not None else None + ), + ) + if env is not None: + env.acquire_generation_lease(host) + env._generation_prepared = True + yield host, binding + if env is not None: + env.release_generation_lease(host) + + +def _candidate(slot, values=(0.0, 0.2, 0.4, 0.6)): + positions = torch.tensor([[value, 0.3] for value in values]).unsqueeze(0) + return CandidateTrajectoryBatch( + positions=positions, + dt=torch.tensor([[0.0] + [0.02] * (len(values) - 1)], dtype=torch.float64), + valid_length=torch.tensor([len(values)]), + identities=( + CandidateIdentity( + f"case-{slot}", + f"initial-{slot}", + f"candidate-{slot}", + "family", + "source", + "rev", + "template", + ), + ), + joint_names=("arm", "locked"), + phases=((TrajectoryPhase("move", 0, len(values), "free"),),), + ) + + +def _executor(host, **kwargs): + kwargs.setdefault("validator", _passed) + kwargs.setdefault("observe", lambda: {"state": host.adapter.robot.qpos}) + return QposRolloutExecutor(host, control_dt=0.02, **kwargs) + + +def _run(executor, binding, candidates, **kwargs): + kwargs.setdefault("on_started", Mock()) + return executor.execute( + binding, + candidates, + { + value.identities[0].candidate_id: (f"episode-{i}", f"commit-{i}") + for i, value in enumerate(candidates) + if value is not None + }, + **kwargs, + ) + + +def test_causal_targets_terminal_frame_ragged_rows_and_owned_payload(scene): + host, binding = scene + robot, sim, env = host.adapter.robot, host.adapter.sim, host.env + if env is not None: + env.post_target_offset = 0.01 + validator = Mock( + side_effect=lambda *args: ( + _passed() if len(sim.updates) == 3 else pytest.fail("early validation") + ) + ) + started = Mock() + episodes = _run( + _executor(host, validator=validator), + binding, + [_candidate(0, (0.0, 0.2)), _candidate(1)], + on_started=started, + ) + assert [len(value.actions) for value in episodes] == [1, 3] + assert started.call_count == 2 + assert sim.updates == [(0.01, 2)] * 3 + assert all(episode.validation.accepted for episode in episodes) + assert episodes[0].observations["state"].shape == (2, 2) + torch.testing.assert_close(episodes[1].actions[:, 0], torch.tensor([0.2, 0.4, 0.6])) + torch.testing.assert_close( + episodes[1].observations["joint_positions"][:, 0], + torch.tensor([0.0, 0.1, 0.25, 0.425]), + ) + torch.testing.assert_close( + episodes[1].timestamps, torch.tensor([0, 0.02, 0.04, 0.06], dtype=torch.float64) + ) + assert robot.commands[2][0, 0] == pytest.approx(0.1) # measured hold + assert robot.target[1, 0] == robot.qpos[1, 0] # final safe hold + if env is not None: + assert env.obs_calls == 4 # initial preparation + actual transitions + assert env._demo_steps.tolist() == [1, 3] + assert [mask.tolist() for mask in env.recorded_masks] == [ + [True, True], + [False, True], + [False, True], + ] + assert env.action_manager.process_action.call_count == 3 + assert env._generation_command_observer is None + robot.qpos.zero_() + assert episodes[1].observations["state"][-1, 0] == pytest.approx(0.425) + + +def test_idle_row_never_starts_or_records(scene): + host, binding = scene + started = Mock() + episodes = _run(_executor(host), binding, [None, _candidate(1)], on_started=started) + assert episodes[0] is None + assert episodes[1].validation.accepted + started.assert_called_once() + assert host.adapter.robot.qpos[0, 0] == 0 + + +@pytest.mark.parametrize("stop_after", [0, 1]) +def test_stop_counts_only_submitted_commands_and_rejects_partial_episode( + scene, stop_after +): + host, binding = scene + started = Mock() + executor = _executor(host) + episodes = _run( + executor, + binding, + [_candidate(0), None], + on_started=started, + should_stop=lambda: len(host.adapter.sim.updates) >= stop_after, + ) + assert started.call_count == int(stop_after > 0) + assert executor.last_failures == { + "candidate-0": "execution cancelled by should_stop" + } + if stop_after == 0: + assert episodes == (None, None) + else: + assert len(episodes[0].actions) == 1 + assert not episodes[0].validation.accepted + assert episodes[0].phases[0].stop_index == 2 + + +@pytest.mark.parametrize("fail_at", [1, 2]) +def test_physics_exception_counts_attempt_without_fabricating_transition( + scene, fail_at +): + host, binding = scene + host.adapter.sim.fail_at = fail_at + started = Mock() + executor = _executor(host) + episodes = _run(executor, binding, [_candidate(0), None], on_started=started) + started.assert_called_once() + assert executor.last_failures == {"candidate-0": "RuntimeError: physics failed"} + if fail_at == 1: + assert episodes[0] is None + else: + assert len(episodes[0].actions) == 1 + assert not episodes[0].validation.accepted + assert host.adapter.robot.target[0, 0] == host.adapter.robot.qpos[0, 0] + + +def test_observation_failure_after_command_still_counts_attempt(scene): + host, binding = scene + if host.env is not None: + host.env.fail_observation_at = 2 + executor = _executor(host) + else: + observe = Mock( + side_effect=[ + {"state": host.adapter.robot.qpos.clone()}, + RuntimeError("camera unavailable"), + ] + ) + executor = _executor(host, observe=observe) + started = Mock() + assert _run(executor, binding, [_candidate(0), None], on_started=started) == ( + None, + None, + ) + started.assert_called_once() + assert executor.last_failures == {"candidate-0": "RuntimeError: camera unavailable"} + + +def test_changed_controller_target_is_recorded_and_rejected(scene): + host, binding = scene + host.adapter.robot.offset = 0.05 + episode = _run(_executor(host), binding, [_candidate(0), None])[0] + torch.testing.assert_close(episode.actions[:, 0], torch.tensor([0.25, 0.45, 0.65])) + assert not episode.validation.accepted + + +def test_changed_obstacle_or_root_never_passes_static_collision_contract(scene): + host, binding = scene + host.adapter.sim.after_update = lambda: host.adapter.sim.cube.pose[0, 0, 3].add_( + 0.01 + ) + episodes = _run(_executor(host), binding, [_candidate(0), _candidate(1)]) + assert not episodes[0].validation.accepted + assert episodes[1].validation.accepted + assert ( + next( + check + for check in episodes[0].validation.checks + if check.check_id == "fixed_collision_world" + ).status + == "failed" + ) + + +def test_actual_clock_disagreement_rejects_without_fabricated_timestamps(scene): + host, binding = scene + host.adapter.sim.clock_scale = 2 + started = Mock() + episodes = _run(_executor(host), binding, [_candidate(0), None], on_started=started) + assert episodes == (None, None) + started.assert_called_once() + + +@pytest.mark.parametrize("fault", ["initial", "inactive", "timing", "case", "budget"]) +def test_preflight_failure_submits_no_command(scene, fault): + host, binding = scene + candidate = _candidate(0) + kwargs = {} + if fault == "initial": + candidate.positions[0, 0, 0] = 0.1 + elif fault == "inactive": + candidate.positions[0, 1, 1] = 0.4 + elif fault == "timing": + candidate.dt[0, 1] = 0.03 + elif fault == "case": + candidate = _candidate(1) + else: + kwargs["max_episode_bytes"] = 1 + with pytest.raises(ValueError): + _run(_executor(host, **kwargs), binding, [candidate, None]) + assert not host.adapter.sim.updates + assert not host.adapter.robot.commands + + +def test_exhaustion_requires_explicit_final_task_success(scene): + host, binding = scene + validator = lambda *args: ValidationResult( + (ValidationCheck("task_success", "failed"),) + ) + episode = _run( + _executor(host, validator=validator), binding, [_candidate(0), None] + )[0] + assert len(episode.actions) == 3 + assert not episode.validation.accepted + + +def test_early_gym_termination_freezes_terminal_observation_and_holds_row(scene): + host, binding = scene + if host.env is None: + pytest.skip("Gym-specific terminal semantics") + host.env.terminate_after[0] = 1 + episodes = _run(_executor(host), binding, [_candidate(0), _candidate(1)]) + assert len(episodes[0].actions) == 1 + assert not episodes[0].validation.accepted + assert len(episodes[1].actions) == 3 + assert episodes[1].validation.accepted + assert episodes[0].observations["state"][-1, 0] == pytest.approx(0.1) + assert host.env._demo_steps.tolist() == [1, 3] + + +def test_schema_growth_is_rejected_before_copying_or_retaining_changed_frame(scene): + host, binding = scene + if host.env is None: + + def observe(): + return { + "state": ( + torch.zeros(2, 3) + if host.adapter.sim.updates + else host.adapter.robot.qpos + ) + } + + executor = _executor(host, observe=observe) + else: + executor = _executor( + host, + encode_observation=lambda raw: { + "state": torch.zeros(2, 3) if host.adapter.sim.updates else raw["state"] + }, + ) + started = Mock() + assert _run(executor, binding, [_candidate(0), None], on_started=started) == ( + None, + None, + ) + started.assert_called_once() + + +def test_oversized_validator_metadata_rejects_with_a_bounded_failure(scene): + host, binding = scene + validator = lambda *args: ValidationResult( + (ValidationCheck("task_success", "passed", "x" * 70000),) + ) + episode = _run( + _executor(host, validator=validator, max_episode_bytes=70000), + binding, + [_candidate(0), None], + )[0] + assert not episode.validation.accepted + assert len(episode.validation.checks[0].detail) < 100 + + +def test_running_guard_covers_final_task_validator(scene): + host, binding = scene + executor = _executor(host) + + def validate(*args): + with pytest.raises(RuntimeError, match="serialized"): + _run(executor, binding, [_candidate(0), None]) + return _passed() + + executor.validator = validate + assert _run(executor, binding, [_candidate(0), None])[0].validation.accepted + + +@pytest.mark.requires_sim +def test_real_cpu_sim_qpos_rollout_records_measured_state_and_time( + tmp_path, monkeypatch +): + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.cfg import RobotCfg, RigidObjectCfg + from embodichain.lab.sim.shapes import CubeCfg + from embodichain.lab.trajectory_generation.initial_state import ( + FixedSceneHost, + InitialStateProfile, + ) + from embodichain.lab.trajectory_generation.integrations.sim import ( + SimInitialStateAdapter, + ) + + monkeypatch.setenv("EMBODICHAIN_SIM_EXIT_PROCESS", "0") + urdf = tmp_path / "execution_robot.urdf" + urdf.write_text( + '' + '' + '' + "" + '' + '' + "" + '' + '' + '', + encoding="utf8", + ) + sim = SimulationManager( + SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=1) + ) + try: + sim.enable_physics(True) + robot = sim.add_robot(RobotCfg(uid="robot", fpath=str(urdf), fix_base=True)) + sim.add_rigid_object( + RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=[0.05, 0.05, 0.05]), + init_pos=[2, 0, 1], + body_type="static", + ) + ) + adapter = SimInitialStateAdapter(sim, robot) + profile = InitialStateProfile( + profile_id="test", + prepare=lambda: None, + signature=lambda: "fixed", + verify=lambda cases: _passed(), + ) + with FixedSceneHost(adapter, profile) as host: + binding = host.acquire_case( + [SceneCase("case", "initial", "fixed", "move", "robot")] + ) + initial = robot.get_qpos().clone() + candidate = CandidateTrajectoryBatch( + positions=torch.stack((initial, initial + 0.01, initial + 0.02), dim=1), + dt=torch.tensor([[0, 0.02, 0.02]], dtype=torch.float64), + valid_length=torch.tensor([3]), + identities=( + CandidateIdentity( + "case", + "initial", + "candidate", + "family", + "source", + "rev", + "template", + ), + ), + joint_names=tuple(robot.joint_names), + ) + executor = QposRolloutExecutor( + host, + control_dt=0.02, + observe=lambda: {"state": robot.get_qpos()}, + validator=_passed, + ) + started = Mock() + episode = executor.execute( + binding, + [candidate], + {"candidate": ("episode", "commit")}, + on_started=started, + )[0] + assert episode is not None + assert episode.validation.accepted, episode.validation.checks + assert episode.actions.shape == (2, 1) + torch.testing.assert_close(episode.actions, candidate.positions[0, 1:]) + torch.testing.assert_close( + episode.observations["joint_positions"][-1], robot.get_qpos()[0] + ) + torch.testing.assert_close( + episode.timestamps, torch.tensor([0, 0.02, 0.04], dtype=torch.float64) + ) + started.assert_called_once() + finally: + sim.destroy() + SimulationManager.flush_cleanup_queue() + + +def test_idle_rows_clear_residual_velocity_targets_and_effort_before_physics(scene): + host, binding = scene + robot = host.adapter.robot + robot.target_velocity.fill_(0.7) + robot.force.fill_(0.8) + + def verify_idle_hold(): + assert robot.target_velocity[0, 0] == 0 + assert robot.force[0, 0] == 0 + + host.adapter.sim.after_update = verify_idle_hold + assert _run(_executor(host), binding, [None, _candidate(1)])[1].validation.accepted + + +def test_started_callback_failure_counts_each_submitted_row_and_stops_before_physics( + scene, +): + host, binding = scene + started = Mock(side_effect=RuntimeError("accounting failed")) + episodes = _run( + _executor(host), binding, [_candidate(0), _candidate(1)], on_started=started + ) + assert episodes == (None, None) + assert started.call_count == 2 + assert not host.adapter.sim.updates + torch.testing.assert_close(host.adapter.robot.target, host.adapter.robot.qpos) + + +def test_measured_inactive_joint_movement_rejects_the_row(scene): + host, binding = scene + host.adapter.sim.after_update = lambda: host.adapter.robot.qpos[0, 1].add_(0.01) + episode = _run(_executor(host), binding, [_candidate(0), None])[0] + assert not episode.validation.accepted + assert "inactive" in episode.validation.checks[0].detail + + +def test_task_validator_cannot_overwrite_execution_check(scene): + host, _ = scene + with pytest.raises(ValueError, match="execution check"): + _executor(host, validator_id="execution_complete") + + +def test_last_failures_are_bounded_read_only_owned_and_cleared_for_the_next_batch( + scene, +): + host, binding = scene + host.adapter.sim.update = Mock(side_effect=RuntimeError("x" * 2000)) + executor = _executor(host) + assert _run(executor, binding, [_candidate(0), None]) == (None, None) + previous = executor.last_failures + assert len(previous["candidate-0"]) == 512 + with pytest.raises(TypeError): + previous["candidate-0"] = "changed" + assert _run(executor, binding, [None, None]) == (None, None) + assert executor.last_failures == {} + assert len(previous["candidate-0"]) == 512 diff --git a/tests/lab/trajectory_generation/test_initial_state_host.py b/tests/lab/trajectory_generation/test_initial_state_host.py new file mode 100644 index 000000000..11d265f6a --- /dev/null +++ b/tests/lab/trajectory_generation/test_initial_state_host.py @@ -0,0 +1,266 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + SceneCase, + ValidationCheck, + ValidationResult, +) +from embodichain.lab.trajectory_generation.initial_state import ( + FixedSceneHost, + InitialStateProfile, +) + + +def _passed(name: str = "task_initial") -> ValidationResult: + return ValidationResult((ValidationCheck(name, "passed"),)) + + +class _Adapter: + def __init__(self) -> None: + self.position = torch.tensor([[1.0], [2.0]]) + self.robot = object() + self.sim = SimpleNamespace(num_envs=2, update=self._settle) + self.settle_calls = 0 + self.restore_calls = 0 + self.settle_delta = 0.0 + self.read_failed = False + + def _settle(self, dt: float, steps: int) -> None: + self.settle_calls += steps + self.position += self.settle_delta + + def signature(self) -> str: + return "two-row-robot" + + def capture(self) -> torch.Tensor: + if self.read_failed: + raise ValueError("backend cannot capture") + return self.position.clone() + + def restore(self, value: torch.Tensor) -> None: + self.restore_calls += 1 + self.position.copy_(value) + + def verify(self, value: torch.Tensor) -> ValidationResult: + return ValidationResult( + ( + ValidationCheck( + "initial_state", + "passed" if torch.equal(self.position, value) else "failed", + ), + ) + ) + + +def _cases() -> tuple[SceneCase, ...]: + return tuple( + SceneCase(f"case-{row}", f"initial-{row}", f"scene-{row}", "pick", "robot") + for row in range(2) + ) + + +def _profile(**kwargs: object) -> InitialStateProfile: + values = dict( + profile_id="fixed_cube", + prepare=lambda: None, + signature=lambda: "fixed-properties", + verify=lambda cases: _passed(), + settling_steps=1, + ) + values.update(kwargs) + return InitialStateProfile(**values) + + +def test_repeated_full_batch_restoration_preserves_case_rows_and_invalidates_tokens() -> ( + None +): + adapter = _Adapter() + with FixedSceneHost(adapter, _profile()) as host: + previous = host.acquire_case(_cases()) + for _ in range(3): + adapter.position += 10 + current = host.restore_initial() + assert torch.equal(adapter.position, torch.tensor([[1.0], [2.0]])) + assert host.verify_initial(current).accepted + assert current.cases == _cases() + assert current.epoch > previous.epoch + with pytest.raises(RuntimeError, match="obsolete"): + host.assert_current(previous) + previous = current + assert adapter.restore_calls == 3 + assert adapter.settle_calls == 4 + with pytest.raises(RuntimeError, match="own"): + host.assert_current(previous) + assert adapter.sim._trajectory_generation_owner is None + + +def test_acquire_captures_after_preparation_and_settling() -> None: + adapter = _Adapter() + adapter.settle_delta = 0.5 + observed = [] + profile = _profile( + prepare=lambda: adapter.position.fill_(3), + verify=lambda cases: (observed.append(adapter.position.clone()) or _passed()), + ) + with FixedSceneHost(adapter, profile) as host: + binding = host.acquire_case(_cases()) + assert torch.equal(observed[0], torch.full((2, 1), 3.5)) + adapter.settle_delta = 0 + adapter.position.fill_(9) + binding = host.restore_initial() + assert torch.equal(adapter.position, torch.full((2, 1), 3.5)) + host.assert_current(binding) + + +def test_fixed_condition_change_rejects_before_prepare_or_restore() -> None: + adapter = _Adapter() + state = {"signature": "initial"} + calls = [] + with FixedSceneHost( + adapter, + _profile( + signature=lambda: state["signature"], + prepare=lambda: calls.append("prepare"), + ), + ) as host: + old = host.acquire_case(_cases()) + state["signature"] = "changed-material" + with pytest.raises(RuntimeError, match="conditions changed"): + host.restore_initial() + assert calls == ["prepare"] + assert adapter.restore_calls == 0 + with pytest.raises(RuntimeError, match="obsolete"): + host.assert_current(old) + + +def test_failed_verification_does_not_publish_binding_and_can_retry_saved_state() -> ( + None +): + adapter = _Adapter() + with FixedSceneHost(adapter, _profile()) as host: + old = host.acquire_case(_cases()) + adapter.settle_delta = 1 + with pytest.raises(RuntimeError, match="verification failed"): + host.restore_initial() + with pytest.raises(RuntimeError, match="obsolete"): + host.assert_current(old) + adapter.settle_delta = 0 + assert host.verify_initial(host.restore_initial()).accepted + + +def test_failed_capture_does_not_create_a_restorable_initial_state() -> None: + adapter = _Adapter() + adapter.read_failed = True + with FixedSceneHost(adapter, _profile()) as host: + with pytest.raises(ValueError, match="cannot capture"): + host.acquire_case(_cases()) + with pytest.raises(RuntimeError, match="acquire_case"): + host.restore_initial() + adapter.read_failed = False + host.assert_current(host.acquire_case(_cases())) + + +def test_second_host_cannot_take_batch_and_closed_host_can_be_replaced() -> None: + adapter = _Adapter() + first = FixedSceneHost(adapter, _profile()) + with pytest.raises(RuntimeError, match="already has"): + FixedSceneHost(adapter, _profile()) + first.close() + first.close() + with FixedSceneHost(adapter, _profile()) as second: + second.assert_current(second.acquire_case(_cases())) + + +def test_partial_batch_is_rejected_without_running_preparation() -> None: + adapter = _Adapter() + calls = [] + with FixedSceneHost(adapter, _profile(prepare=lambda: calls.append(True))) as host: + with pytest.raises(ValueError, match="every simulator row"): + host.acquire_case(_cases()[:1]) + assert not calls + host.acquire_case(_cases()) + with pytest.raises(RuntimeError, match="already acquired"): + host.acquire_case(_cases()) + + +@pytest.mark.parametrize("status", ["failed", "unavailable", "not_run"]) +def test_profile_must_run_and_pass_all_required_checks(status: str) -> None: + adapter = _Adapter() + profile = _profile( + verify=lambda cases: ValidationResult( + (ValidationCheck("scene_signature", status, "cannot certify"),) + ) + ) + with FixedSceneHost(adapter, profile) as host: + with pytest.raises(RuntimeError, match="scene_signature"): + host.acquire_case(_cases()) + + +@pytest.mark.parametrize( + "override", + [ + {"physics_dt": 0}, + {"physics_dt": float("nan")}, + {"settling_steps": -1}, + {"settling_steps": True}, + {"prepare": None}, + {"profile_id": ""}, + {"allowed_interval_events": ("push", "push")}, + ], +) +def test_invalid_preparation_profile_is_rejected(override: dict) -> None: + with pytest.raises(ValueError): + _profile(**override) + + +def test_gym_unknown_interval_event_is_rejected_before_acquiring_lease() -> None: + adapter = _Adapter() + env = SimpleNamespace( + sim=adapter.sim, + robot=adapter.robot, + num_envs=2, + physics_dt=0.01, + event_manager=SimpleNamespace(active_functors={"interval": ["random_push"]}), + ) + with pytest.raises(ValueError, match="uncertified interval"): + FixedSceneHost(adapter, _profile(), env=env) + assert not hasattr(adapter.sim, "_trajectory_generation_owner") + + +def test_normal_sim_reset_is_blocked_until_host_releases_batch() -> None: + from embodichain.lab.sim.sim_manager import SimulationManager as SimManager + + adapter = _Adapter() + with FixedSceneHost(adapter, _profile()): + with pytest.raises(RuntimeError, match="generation host owns"): + SimManager.reset_objects_state(adapter.sim) + + +def test_topology_change_invalidates_prepared_execution() -> None: + adapter = _Adapter() + with FixedSceneHost(adapter, _profile()) as host: + binding = host.acquire_case(_cases()) + adapter.signature = lambda: "new-robot-control-layout" + with pytest.raises(RuntimeError, match="topology"): + host.assert_current(binding) diff --git a/tests/lab/trajectory_generation/test_pickup_collection.py b/tests/lab/trajectory_generation/test_pickup_collection.py new file mode 100644 index 000000000..93e24e6ea --- /dev/null +++ b/tests/lab/trajectory_generation/test_pickup_collection.py @@ -0,0 +1,176 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""End-to-end parallel PickUp collection with physical contacts and sealed data.""" + +from __future__ import annotations + +import json +from pathlib import Path +import subprocess +import sys + +import numpy as np +import pytest + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.requires_sim +@pytest.mark.subprocess_sim +def test_pickup_collection(tmp_path: Path) -> None: + av = pytest.importorskip("av") + import pyarrow.parquet as pq + + root = Path(__file__).resolve().parents[3] + output = tmp_path / "collection" + process = subprocess.run( + [ + sys.executable, + "examples/sim/motion/trajectory_generation/cube_pickup_collection.py", + "--output", + str(output), + "--episodes", + "8", + "--record-video", + ], + cwd=root, + capture_output=True, + text=True, + timeout=480, + ) + assert process.returncode == 0, process.stdout[-4000:] + process.stderr[-4000:] + report = json.loads((output / "generation_report.json").read_text()) + pickup = json.loads((output / "pickup_report.json").read_text()) + manifest = json.loads((output / "manifest.json").read_text()) + assert report["target_reached"] and report["error"] is None + assert report["counts"]["committed"] == len(manifest["episodes"]) == 8 + assert pickup["prepared_batches"] >= 2 + assert report["configuration"]["source"]["kind"] == "atomic" + initial_joints, initial_objects, cases, closing_axes, paths = [], [], set(), {}, {} + per_case_initial = {} + for record in manifest["episodes"]: + shard = output / record["shard"] + evidence = json.loads((shard / "episode.json").read_text()) + case = evidence["identity"]["scene_case_id"] + cases.add(case) + checks = {check["check_id"]: check for check in evidence["validation"]} + assert all(check["status"] == "passed" for check in checks.values()) + assert { + "planned_path_collision", + "path_collision", + "physical_contacts", + "held_object_stability", + "actual_motion_limits", + "task_success", + "fixed_collision_world", + } <= checks.keys() + metrics = checks["held_object_stability"]["metrics"] + assert metrics["minimum_lift_m"] >= 0.12 + assert ( + min( + metrics["finger_0_contact_fraction"], + metrics["finger_1_contact_fraction"], + ) + >= 0.95 + ) + assert metrics["max_relative_translation_m"] <= 0.01 + assert metrics["hold_seconds"] >= 1.0 + assert checks["physical_contacts"]["metrics"]["physics_samples"] == 2710 + assert evidence["observation_shapes"]["object_pose"] == [4, 4] + assert len(evidence["metadata"]["commanded_joint_indices"]) == 7 + assert len(evidence["metadata"]["joint_names"]) == 8 + files = sorted((shard / "dataset" / "data").rglob("*.parquet")) + table = pq.read_table(files) + assert table.num_rows == 271 + joints = np.array(table["observation.joint_positions"].to_pylist()) + objects = np.array(table["observation.object_pose"].to_pylist()).reshape( + -1, 4, 4 + ) + tcp = np.array(table["observation.tcp_pose"].to_pylist()).reshape(-1, 4, 4) + initial_joints.append(joints[0]) + initial_objects.append(objects[0]) + if case in per_case_initial: + np.testing.assert_allclose( + objects[0], per_case_initial[case], atol=1e-6, rtol=0 + ) + else: + per_case_initial[case] = objects[0] + paths.setdefault(case, []).append(tcp[:81, :3, 3]) + with np.load(shard / "terminal.npz", allow_pickle=False) as terminal: + np.testing.assert_allclose( + np.diff(terminal["timestamps"]), 0.05, atol=1e-8, rtol=0 + ) + assert terminal["timestamps"].shape == (272,) + assert terminal["observation.object_pose"][2, 3] - objects[0, 2, 3] >= 0.12 + closing_axes[case] = terminal["observation.tcp_pose"][:3, 0] + assert cases == {f"cube_grasp_row_{row}" for row in range(4)} + np.testing.assert_allclose( + initial_joints, np.broadcast_to(initial_joints[0], (8, 8)), atol=1e-6, rtol=0 + ) + np.testing.assert_allclose( + initial_objects, + np.broadcast_to(initial_objects[0], (8, 4, 4)), + atol=1e-5, + rtol=0, + ) + assert ( + abs(np.dot(closing_axes["cube_grasp_row_0"], closing_axes["cube_grasp_row_2"])) + < 0.1 + ) + assert any( + len(values) >= 2 and np.linalg.norm(values[0] - values[1], axis=-1).max() > 0.01 + for values in paths.values() + ) + with av.open(str(output / "preview.mp4")) as video: + stream = video.streams.video[0] + assert stream.codec_context.name == "h264" + assert (stream.width, stream.height) == (1280, 1056) + assert float(stream.average_rate) == 20.0 + frames = sum(1 for _ in video.decode(stream)) + assert frames == pickup["video_frames"] == pickup["prepared_batches"] * 272 + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.requires_sim +@pytest.mark.subprocess_sim +def test_pickup_cli_releases_native_scene_on_integration_failure( + tmp_path: Path, +) -> None: + root = Path(__file__).resolve().parents[3] + script = """ +import sys +from examples.sim.motion.trajectory_generation.cube_pickup_collection import main +from embodichain.lab.trajectory_generation.integrations.contact import PickUpMotionValidator + +def refuse(self, *args, **kwargs): + raise ValueError("injected contact backend failure") + +PickUpMotionValidator.__init__ = refuse +sys.argv = ["cube_pickup_collection.py", "--output", sys.argv[1], "--episodes", "1"] +main() +""" + process = subprocess.run( + [sys.executable, "-c", script, str(tmp_path / "failed")], + cwd=root, + capture_output=True, + text=True, + timeout=90, + ) + assert process.returncode == 1, process.stdout[-2000:] + process.stderr[-2000:] + assert "injected contact backend failure" in process.stderr + assert not (tmp_path / "failed" / "manifest.json").exists() diff --git a/tests/lab/trajectory_generation/test_planning.py b/tests/lab/trajectory_generation/test_planning.py new file mode 100644 index 000000000..1f389ca66 --- /dev/null +++ b/tests/lab/trajectory_generation/test_planning.py @@ -0,0 +1,570 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Real-batch mapping and failure isolation for the free-motion adapter.""" + +from __future__ import annotations + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.lab.sim.motion.planners import CollisionWorldInfo, MotionGenerator +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CandidateIdentity, + CandidateTrajectoryBatch, + MotionSnapshot, + SceneCase, + TrajectoryPhase, +) +from embodichain.lab.trajectory_generation.integrations.planning import ( + EEFPath, + EnvRowMotionPlanner, +) + + +@pytest.mark.gpu +@pytest.mark.slow +def test_real_curobo_free_path_and_dynamic_obstacle_with_cpu_physics(): + """Exercise the actual joint-bound, self, and world checker on one Panda.""" + pytest.importorskip("curobo") + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.cfg import RigidBodyAttributesCfg + from embodichain.lab.sim.motion.planners import MotionGenCfg + from embodichain.lab.sim.motion.planners.curobo.curobo_planner import ( + CuroboPlannerCfg, + CuroboWorldCfg, + ) + from embodichain.lab.sim.objects import RigidObjectCfg + from embodichain.lab.sim.robots import FrankaPandaCfg + from embodichain.lab.sim.shapes import CubeCfg + + sim = SimulationManager( + SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=1) + ) + generator = None + try: + robot = sim.add_robot( + cfg=FrankaPandaCfg.from_dict( + {"uid": "free_path_panda", "robot_type": "panda"} + ) + ) + robot.set_qpos( + torch.tensor([robot.cfg.init_qpos], device=robot.device), target=False + ) + obstacle = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="path_obstacle", + shape=CubeCfg(size=(0.3, 0.3, 0.3)), + attrs=RigidBodyAttributesCfg(), + body_type="kinematic", + init_pos=(3.0, 0.0, 1.0), + ) + ) + generator = MotionGenerator( + MotionGenCfg( + planner_cfg=CuroboPlannerCfg( + robot_uid=robot.uid, + world=CuroboWorldCfg( + rigid_objects={"obstacle": obstacle}, + dynamic_obstacle_names=["obstacle"], + obstacle_representation="cuboid", + ), + use_cuda_graph=False, + cuda_device=0, + ) + ) + ) + planner = EnvRowMotionPlanner( + generator, + control_part="arm", + held_object_ids=(), + max_joint_step=0.01, + max_validation_samples=8, + ) + snapshot = MotionSnapshot( + SceneCase("real-case", "real-initial", "real-scene", "move", "panda"), + tuple(robot.joint_names), + robot.get_qpos()[0], + robot.get_qvel()[0], + robot.get_local_pose(to_matrix=True)[0], + {"obstacle": obstacle.get_local_pose(to_matrix=True)[0]}, + ) + positions = snapshot.joint_positions.repeat(1, 3, 1) + positions[0, :, planner.joint_ids[0]] += torch.tensor([0.0, 0.01, 0.02]) + identity = CandidateIdentity( + "real-case", + "real-initial", + "real-candidate", + "real-family", + "source", + "r1", + "template", + ) + batch = CandidateTrajectoryBatch( + positions, + torch.tensor([[0.0, 0.1, 0.1]]), + torch.tensor([3]), + (identity,), + snapshot.joint_names, + ((TrajectoryPhase("free", 0, 3),),), + source_row_indices=torch.tensor([0]), + ) + result = planner.validate_qpos(batch, (snapshot,))[0] + assert result.accepted, ( + result.checks, + snapshot.joint_positions, + robot.cfg.init_qpos, + planner.fixed_ids, + ) + assert robot.device.type == "cpu" + assert generator.planner._curobo_device.type == "cuda" + assert len(generator.planner._backend_cache) == 1 + + blocked_pose = robot.compute_fk( + qpos=robot.get_qpos(name="arm"), name="arm", to_matrix=True + ) + obstacle.set_local_pose(blocked_pose) + blocked = replace(snapshot, entity_poses={"obstacle": blocked_pose[0]}) + rejected = planner.validate_qpos(batch, (blocked,))[0] + assert rejected.checks[0].status == "failed", rejected.checks + assert len(generator.planner._backend_cache) == 1 + finally: + if generator is not None: + generator.planner.close() + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + +class _Robot: + device = torch.device("cpu") + num_instances = 4 + joint_names = ("arm", "gripper") + cfg = SimpleNamespace(init_qpos=[0.0, 0.25]) + + def __init__(self): + self.roots = torch.eye(4).repeat(self.num_instances, 1, 1) + self.roots[:, 0, 3] = torch.arange(self.num_instances) * 10.0 + self.ik_calls = [] + self.fk_calls = [] + self.failed_ik_rows = set() + + def get_joint_ids(self, name): + assert name == "arm" + return [0] + + def get_local_pose(self, *, to_matrix): + assert to_matrix + return self.roots.clone() + + def get_qpos_limits(self): + return torch.tensor([[-10.0, 10.0], [0.0, 1.0]]).repeat( + self.num_instances, 1, 1 + ) + + def compute_fk(self, *, qpos, name, to_matrix): + assert qpos.shape == (self.num_instances, 1) + assert torch.isfinite(qpos).all() + self.fk_calls.append(qpos.clone()) + poses = self.roots.clone() + poses[:, 0, 3] += qpos[:, 0] + return poses + + def compute_ik(self, *, pose, name, joint_seed): + assert pose.shape == (self.num_instances, 4, 4) + assert torch.isfinite(joint_seed).all() + self.ik_calls.append((pose.clone(), joint_seed.clone())) + solved = (pose[:, 0, 3] - self.roots[:, 0, 3]).unsqueeze(1) + mask = torch.ones(self.num_instances, dtype=torch.bool) + for row in self.failed_ik_rows: + solved[row] = float("nan") + mask[row] = False + return mask, solved + + +class _Backend: + supports_joint_trajectory_validation = True + collision_world_info = CollisionWorldInfo( + ("obstacle",), ("obstacle",), "per_env", True + ) + + def __init__(self): + self.queries = [] + self.block_middle = False + self.reject_padding_after = {} + + def validate_joint_trajectory(self, trajectory, *, control_part, obstacle_poses): + assert trajectory.shape[0] == 4 + assert torch.isfinite(trajectory).all() + self.queries.append((trajectory.clone(), obstacle_poses)) + mask = torch.ones(trajectory.shape[:2], dtype=torch.bool) + if self.block_middle: + mask &= ~((trajectory[..., 0] > 0.04) & (trajectory[..., 0] < 0.06)) + for row, start in self.reject_padding_after.items(): + mask[row, start:] = False + return mask + + +def _fixture(**kwargs): + robot, backend = _Robot(), _Backend() + generator = object.__new__(MotionGenerator) + generator.robot = robot + generator.planner = backend + planner = EnvRowMotionPlanner( + generator, control_part="arm", held_object_ids=(), **kwargs + ) + snapshots = tuple( + MotionSnapshot( + SceneCase(f"case-{row}", f"initial-{row}", f"scene-{row}", "move", "robot"), + robot.joint_names, + torch.tensor([0.0, 0.25]), + torch.zeros(2), + robot.roots[row], + {"obstacle": robot.roots[row]}, + ) + for row in range(robot.num_instances) + ) + return planner, robot, backend, snapshots + + +def _identity(index, row): + return CandidateIdentity( + f"case-{row}", + f"initial-{row}", + f"candidate-{index}", + f"family-{index}", + "source", + "revision", + "template", + ) + + +def _batch(rows, *, displacements=None, phases=None): + count = len(rows) + values = [0.1] * count if displacements is None else displacements + q = torch.tensor([[[0.0, 0.25], [value, 0.25]] for value in values]).reshape( + count, 2, 2 + ) + return CandidateTrajectoryBatch( + q, + torch.tensor([[0.0, 0.1]] * count).reshape(count, 2), + torch.full((count,), 2, dtype=torch.int64), + tuple(_identity(index, row) for index, row in enumerate(rows)), + ("arm", "gripper"), + ( + tuple((TrajectoryPhase("free", 0, 2),) for _ in rows) + if phases is None + else phases + ), + source_row_indices=torch.tensor(rows, dtype=torch.int64), + ) + + +def _eef_paths(rows, robot): + paths = [] + for index, row in enumerate(rows): + poses = robot.roots[row].repeat(3, 1, 1) + poses[:, 0, 3] += torch.tensor([0.0, 0.02, 0.04 + index * 0.003]) + paths.append( + EEFPath( + _identity(index, row), + row, + poses, + torch.tensor([0.0, 0.1, 0.2]), + (TrajectoryPhase("free", 0, 3),), + ) + ) + return paths + + +def test_eleven_candidates_use_four_real_rows_across_three_rounds(): + planner, robot, backend, snapshots = _fixture() + rows = [0, 0, 1, 2, 3, 3, 3, 1, 2, 0, 2] + batch = _batch(rows, displacements=[0.03 + index * 0.001 for index in range(11)]) + original = batch.positions.clone() + results = planner.validate_qpos(batch, snapshots) + assert all(result.accepted for result in results) + assert len(backend.queries) == 3 + assert robot.num_instances == 4 + assert torch.equal(batch.positions, original) + expected_groups = [ + {0: 0, 1: 2, 2: 3, 3: 4}, + {0: 1, 1: 7, 2: 8, 3: 5}, + {0: 9, 2: 10, 3: 6}, + ] + for (query, obstacles), expected in zip(backend.queries, expected_groups): + assert query.shape[0] == 4 + assert torch.equal(obstacles["obstacle"], robot.roots) + for row in range(4): + assert query[row, -1, 0] == pytest.approx( + batch.positions[expected[row], -1, 0].item() if row in expected else 0.0 + ) + + +def test_eef_mapping_retains_all_rows_times_phases_and_uncontrolled_joints(): + planner, robot, backend, snapshots = _fixture() + rows = [0, 0, 1, 2, 3, 3, 3, 1, 2, 0, 2] + paths = _eef_paths(rows, robot) + batch, results = planner.plan_eef(paths, snapshots) + assert all(result.accepted for result in results) + assert len(robot.ik_calls) == 6 + assert batch.source_row_indices.tolist() == rows + for index, path in enumerate(paths): + assert torch.allclose( + batch.positions[index, :, 0], + path.poses[:, 0, 3] - robot.roots[rows[index], 0, 3], + ) + assert torch.equal(batch.dt[index], path.dt) + assert batch.phases[index] == path.phases + assert torch.equal(batch.positions[..., 1], torch.full((11, 3), 0.25)) + + +def test_fixed_solved_joint_samples_skip_ik_and_remain_exact(): + planner, robot, backend, snapshots = _fixture() + path = _eef_paths([2], robot)[0] + solved = torch.tensor([[0.0], [0.02], [0.04]]) + path = replace(path, solved_joint_targets=solved) + batch, results = planner.plan_eef([path], snapshots) + assert results[0].accepted + assert not robot.ik_calls + assert torch.equal(batch.positions[0, :, 0], solved[:, 0]) + bad_path = replace(path, solved_joint_targets=solved + 0.1) + _, rejected = planner.plan_eef([bad_path], snapshots) + assert not rejected[0].accepted + assert not robot.ik_calls + + +def test_failed_ik_row_never_contaminates_later_seeds_or_other_rows(): + planner, robot, backend, snapshots = _fixture() + robot.failed_ik_rows = {1} + batch, results = planner.plan_eef(_eef_paths([0, 1, 2], robot), snapshots) + assert [value.accepted for value in results] == [True, False, True] + assert torch.equal(batch.positions[1], snapshots[1].joint_positions.expand(3, -1)) + assert all(torch.isfinite(seed).all() for _, seed in robot.ik_calls) + + +def test_collision_between_anchors_and_across_phase_boundary_is_rejected(): + planner, robot, backend, snapshots = _fixture(max_joint_step=0.01) + backend.block_middle = True + phases = ((TrajectoryPhase("first", 0, 1), TrajectoryPhase("second", 1, 2)),) + result = planner.validate_qpos(_batch([0], phases=phases), snapshots)[0] + assert result.checks[0].status == "failed" + assert backend.queries[0][0].shape[1] > 2 + + +def test_inactive_rows_and_padded_samples_do_not_affect_acceptance(): + planner, robot, backend, snapshots = _fixture(max_joint_step=0.02) + backend.reject_padding_after = {0: 2, 2: 0, 3: 0} + results = planner.validate_qpos( + _batch([0, 1], displacements=[0.01, 0.1]), snapshots + ) + assert all(value.accepted for value in results) + query = backend.queries[0][0] + assert torch.equal(query[2:], torch.zeros_like(query[2:])) + + +@pytest.mark.parametrize("kind", ["contact", "hold", "missing"]) +def test_unsupported_phase_semantics_never_enter_collision_backend(kind): + planner, robot, backend, snapshots = _fixture() + phases = ((),) if kind == "missing" else ((TrajectoryPhase("phase", 0, 2, kind),),) + result = planner.validate_qpos(_batch([0], phases=phases), snapshots)[0] + assert result.checks[0].status == "unavailable" + assert not backend.queries + + +def test_held_objects_and_moving_or_noninitial_grippers_are_unavailable(): + planner, robot, backend, snapshots = _fixture() + batch = _batch([0]) + planner.held_object_ids = ("payload",) + assert planner.validate_qpos(batch, snapshots)[0].checks[0].status == "unavailable" + planner.held_object_ids = () + batch.positions[0, 1, 1] = 0.4 + assert planner.validate_qpos(batch, snapshots)[0].checks[0].status == "unavailable" + batch.positions[0, :, 1] = 0.4 + snapshots = ( + replace(snapshots[0], joint_positions=torch.tensor([0.0, 0.4])), + *snapshots[1:], + ) + assert planner.validate_qpos(batch, snapshots)[0].checks[0].status == "unavailable" + assert not backend.queries + + +def test_sampling_budget_and_missing_backend_capability_are_unavailable(): + planner, robot, backend, snapshots = _fixture(max_validation_samples=2) + assert ( + planner.validate_qpos(_batch([0]), snapshots)[0].checks[0].status + == "unavailable" + ) + backend.supports_joint_trajectory_validation = False + assert ( + planner.validate_qpos(_batch([0]), snapshots)[0].checks[0].status + == "unavailable" + ) + assert not backend.queries + + +def test_empty_qpos_and_eef_inputs_do_not_invoke_backend(): + planner, robot, backend, snapshots = _fixture() + assert planner.validate_qpos(_batch([]), snapshots) == () + batch, results = planner.plan_eef([], snapshots) + assert batch.positions.shape == (0, 0, 2) + assert results == () + assert not backend.queries and not robot.ik_calls + + +def test_nan_candidate_is_failed_before_backend_and_eef_constructor_rejects_nan(): + planner, robot, backend, snapshots = _fixture() + batch = _batch([0]) + batch.positions[0, 1, 0] = float("nan") + assert planner.validate_qpos(batch, snapshots)[0].checks[0].status == "failed" + path = _eef_paths([0], robot)[0] + poses = path.poses.clone() + poses[1, 0, 0] = float("nan") + with pytest.raises(ValueError, match="finite"): + replace(path, poses=poses) + assert not backend.queries + + +def test_missing_source_mapping_foreign_case_and_stale_roots_are_rejected(): + planner, robot, backend, snapshots = _fixture() + with pytest.raises(ValueError, match="source_row_indices"): + planner.validate_qpos(replace(_batch([0]), source_row_indices=None), snapshots) + with pytest.raises(ValueError, match="identity"): + planner.validate_qpos( + replace(_batch([0]), source_row_indices=torch.tensor([1])), snapshots + ) + robot.roots[0, 0, 3] += 1 + with pytest.raises(ValueError, match="root pose"): + planner.validate_qpos(_batch([0]), snapshots) + assert not backend.queries + + +def test_collision_densification_is_interleaved_with_real_batch_queries(monkeypatch): + """The adapter never materializes dense samples for every logical candidate.""" + planner, robot, backend, snapshots = _fixture() + batch = _batch([0, 0, 1, 2, 3, 3, 3, 1, 2, 0, 2]) + events = [] + real_arange = torch.arange + real_validate = backend.validate_joint_trajectory + + def arange(*args, **kwargs): + events.append("densify") + return real_arange(*args, **kwargs) + + def validate(*args, **kwargs): + events.append("query") + return real_validate(*args, **kwargs) + + monkeypatch.setattr(torch, "arange", arange) + monkeypatch.setattr(backend, "validate_joint_trajectory", validate) + assert all(result.accepted for result in planner.validate_qpos(batch, snapshots)) + dense_since_query = 0 + for event in events: + if event == "densify": + dense_since_query += 1 + assert dense_since_query <= robot.num_instances + else: + dense_since_query = 0 + assert events[-1] == "query" + + +def test_eef_keeps_float64_arrival_intervals_and_solved_joint_samples(): + planner, robot, backend, snapshots = _fixture() + path = _eef_paths([0], robot)[0] + dt = torch.tensor([0.0, 0.10000000001, 0.20000000003], dtype=torch.float64) + solved = torch.tensor( + [[0.0], [0.02000000001], [0.04000000003]], dtype=torch.float64 + ) + batch, results = planner.plan_eef( + [replace(path, dt=dt, solved_joint_targets=solved)], snapshots + ) + assert results[0].accepted + assert batch.dt.dtype == torch.float64 + assert torch.equal(batch.dt[0], dt) + assert torch.equal(batch.positions[0, :, 0], solved[:, 0]) + + +def test_missing_locked_joint_configuration_is_unavailable(): + planner, robot, backend, snapshots = _fixture() + robot.cfg = SimpleNamespace(init_qpos=None) + result = planner.validate_qpos(_batch([0]), snapshots)[0] + assert result.checks[0].status == "unavailable" + assert not backend.queries + + +def test_post_rollout_qpos_validation_uses_snapshots_without_reading_current_joints(): + planner, robot, backend, snapshots = _fixture() + + def forbid_live_qpos(*args, **kwargs): + raise AssertionError("post-rollout current joints do not represent the path") + + robot.get_qpos = forbid_live_qpos + measured = _batch([0]) + result = planner.validate_qpos(measured, snapshots)[0] + assert result.accepted + query, obstacles = backend.queries[0] + assert query[0, -1, 0] == measured.positions[0, -1, 0] + assert torch.equal(obstacles["obstacle"], robot.roots) + + +@pytest.mark.parametrize( + "start_error, accepted", [(0.0, True), (5e-7, True), (5e-6, False)] +) +def test_measured_initial_anchor_has_a_stricter_tolerance_than_host_default( + start_error, accepted +): + planner, robot, backend, snapshots = _fixture() + measured = _batch([0]) + measured.positions[0, 0, 0] += start_error + result = planner.validate_qpos(measured, snapshots)[0] + assert result.accepted is accepted + if not accepted: + assert result.checks[0].status == "failed" + assert "does not start" in result.checks[0].detail + assert not backend.queries + + +def test_changed_runtime_lock_model_is_unavailable_even_for_matching_new_snapshot(): + from embodichain.lab.sim.motion.planners.curobo.curobo_planner import ( + CuroboPlanner, + CuroboPlannerCfg, + ) + from embodichain.lab.sim.motion.planners import MoveType + + planner, robot, _, snapshots = _fixture() + robot.control_parts = {"arm": ["arm"]} + curobo = object.__new__(CuroboPlanner) + curobo.robot = robot + curobo.cfg = CuroboPlannerCfg(robot_uid="cached-model") + curobo._curobo_device = torch.device("cpu") + curobo._backend_cache = { + ("arm", 4, False, MoveType.JOINT_MOVE): SimpleNamespace( + robot_lock_signature=curobo._locked_joint_signature("arm") + ) + } + planner.motion_generator.planner = curobo + robot.cfg = SimpleNamespace(init_qpos=[0.0, 0.3]) + snapshots = tuple( + replace(snapshot, joint_positions=torch.tensor([0.0, 0.3])) + for snapshot in snapshots + ) + candidate = _batch([0]) + candidate.positions[:, :, 1] = 0.3 + result = planner.validate_qpos(candidate, snapshots)[0] + assert result.checks[0].status == "unavailable" + assert "locked-joint configuration changed" in result.checks[0].detail diff --git a/tests/lab/trajectory_generation/test_runner.py b/tests/lab/trajectory_generation/test_runner.py new file mode 100644 index 000000000..3652dc87a --- /dev/null +++ b/tests/lab/trajectory_generation/test_runner.py @@ -0,0 +1,506 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Runner scheduling/accounting with tensor hosts and real local LeRobot writes. + +These tests isolate orchestration; the tensor executor does not certify physics. +""" + +from __future__ import annotations + +from dataclasses import replace +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + ExpertEpisode, + MotionSnapshot, + SceneCase, + TrajectoryGenerationJobCfg, + TrajectoryPhase, + TrajectoryTemplate, + ValidationCheck, + ValidationResult, +) +from embodichain.lab.trajectory_generation.runner import ( + GenerationRunner, + MotionLimitsProfile, +) +from embodichain.lab.trajectory_generation.sinks import LeRobotEpisodeSink + + +def _checks(name="task_success", status="passed"): + return ValidationResult((ValidationCheck(name, status),)) + + +class _Host: + def __init__(self, rows): + self.robot = SimpleNamespace( + joint_names=("a", "b"), + get_qpos_limits=lambda: torch.tensor([[-2.0, 2.0], [-2.0, 2.0]]).repeat( + rows, 1, 1 + ), + ) + self.adapter = SimpleNamespace( + robot=self.robot, + sim=SimpleNamespace(num_envs=rows), + tolerances_profile_id="fixed_scene_tolerances", + ) + self.profile = SimpleNamespace(profile_id="fixed_scene_initial_state") + self.prepared = 0 + self.closed = False + self.cases = () + self.fail_restore = False + + def acquire_case(self, cases): + self.cases = cases + self.prepared += 1 + return self.prepared + + def restore_initial(self): + self.prepared += 1 + if self.fail_restore: + raise RuntimeError("restoration mismatch") + return self.prepared + + def snapshots(self, binding): + return tuple( + MotionSnapshot( + case, + self.robot.joint_names, + torch.zeros(2), + torch.zeros(2), + torch.eye(4), + ) + for case in self.cases + ) + + def verify_initial(self, binding): + return _checks("initial_state") + + def close(self): + self.closed = True + + +class _Planner: + def __init__(self, host): + self.robot = host.robot + self.motion_generator = SimpleNamespace( + collision_world_entity_ids=(), dynamic_collision_entity_ids=() + ) + self.calls = [] + self.planning_status = "passed" + self.actual_status = "passed" + self.executed = False + self.reject_first = 0 + + def validate_qpos(self, batch, snapshots): + self.calls.append(batch) + status = self.actual_status if self.executed else self.planning_status + if len(self.calls) <= self.reject_first: + status = "failed" + return (_checks("path_collision", status),) + + +class _Executor: + control_dt = 0.05 + validator_id = "task_success" + validation_profile_id = "verified_motion" + max_episode_bytes = 16384 + + def __init__(self, host, planner): + self.host, self.planner = host, planner + self.calls = [] + self.episodes = [] + self.fail_before = False + self.fail_after = False + self.status = "passed" + self.actual_middle = None + self.unsupported_observation = False + self.after_first = lambda: None + self.last_failures = {} + self.no_transition_error = None + + def execute(self, binding, candidates, episode_ids, *, on_started, should_stop): + self.calls.append(tuple(candidate is not None for candidate in candidates)) + if self.fail_before: + raise RuntimeError("before command") + episodes = [] + for candidate in candidates: + if candidate is None or should_stop(): + episodes.append(None) + continue + identity = candidate.identities[0] + on_started(identity) + if self.no_transition_error: + self.last_failures[identity.candidate_id] = self.no_transition_error + episodes.append(None) + continue + if self.fail_after: + raise RuntimeError("after command") + count = int(candidate.valid_length[0]) + actual = candidate.positions[0, :count].clone() + if self.actual_middle is not None: + actual[1] = self.actual_middle + observations = {"joint_positions": actual} + if self.unsupported_observation: + observations["unsupported"] = torch.ones(count, 2, dtype=torch.bool) + episode = ExpertEpisode( + identity, + observations, + candidate.positions[0, 1:count], + torch.arange(count, dtype=torch.float64) * self.control_dt, + "qpos", + ValidationResult( + ( + ValidationCheck("execution", "passed"), + ValidationCheck("fixed_collision_world", "passed"), + ValidationCheck("task_success", self.status), + ) + ), + *episode_ids[identity.candidate_id], + phases=candidate.phases[0], + ) + episodes.append(episode) + self.episodes.append(episode) + self.after_first() + self.planner.executed = True + return tuple(episodes) + + +def _job(tmp_path, *, rows=1, configuration=None): + cfg = TrajectoryGenerationJobCfg.from_mapping(configuration or {}) + host = _Host(rows) + planner = _Planner(host) + executor = _Executor(host, planner) + sink = LeRobotEpisodeSink(tmp_path / "collection", fps=20, max_episode_bytes=16384) + runner = GenerationRunner( + cfg, + host, + planner, + executor, + sink, + motion_limits=MotionLimitsProfile( + torch.full((2,), 100.0), torch.full((2,), 10000.0) + ), + ) + cases = tuple( + SceneCase(f"case-{row}", f"initial-{row}", "scene", "move", "robot") + for row in range(rows) + ) + templates = tuple( + TrajectoryTemplate( + "handwritten_qpos", + "v1", + "reference_0", + host.robot.joint_names, + torch.tensor([[0.0, 0.0], [0.2, 0.1], [0.4, 0.0]]), + torch.tensor([0.0, 0.05, 0.05], dtype=torch.float64), + ( + TrajectoryPhase( + "free", 0, 3, allowed_operators=("joint_residual", "retime") + ), + ), + ("joint_residual", "retime"), + validator_id="task_success", + controlled_joint_indices=(0, 1), + ) + for _ in range(rows) + ) + return runner, cases, templates + + +def _report(runner): + return json.loads((runner.sink.root / "generation_report.json").read_text()) + + +def test_multiple_rounds_and_tail_commit_only_reserved_actual_evidence(tmp_path): + runner, cases, templates = _job( + tmp_path, + rows=2, + configuration={ + "augmentation": { + "factors": {"spatial": {"enabled": True, "joint_offset_scale": 0.01}}, + "coverage": {"joint_dedup_normalized_tol": 0.000001}, + }, + "collection": {"target_committed_episodes": 3, "max_proposals": 8}, + }, + ) + report = runner.run(cases, templates) + assert report["target_reached"] + assert report["counts"]["committed"] == report["counts"]["rollout_attempted"] == 3 + assert runner.executor.calls == [(True, True), (False, True)] + assert runner.host.prepared == 2 + assert runner.host.closed and not (runner.sink.root / ".writer.lock").exists() + assert report["pending_reserved_bytes"] == 0 + manifest = json.loads((runner.sink.root / "manifest.json").read_text()) + assert len(manifest["episodes"]) == 3 + for record in manifest["episodes"]: + evidence = json.loads( + (runner.sink.root / record["shard"] / "episode.json").read_text() + ) + checks = {check["check_id"] for check in evidence["validation"]} + assert { + "planned_path_collision", + "path_collision", + "actual_motion_limits", + "motion_quality", + "task_success", + } <= checks + assert _report(runner)["counts"]["committed"] == 3 + with pytest.raises(RuntimeError, match="single-use"): + runner.run(cases, templates) + + +def test_planning_failures_do_not_restore_or_count_rollouts(tmp_path): + runner, cases, templates = _job(tmp_path) + runner.planner.reject_first = 2 + report = runner.run(cases, templates) + assert report["counts"]["proposed"] == 3 + assert report["counts"]["rollout_attempted"] == 1 + assert report["counts"]["committed"] == 1 + assert runner.host.prepared == 1 + + +def test_required_backend_unavailable_fails_immediately_with_audit(tmp_path): + runner, cases, templates = _job(tmp_path) + runner.planner.planning_status = "unavailable" + with pytest.raises(RuntimeError, match="capability is unavailable"): + runner.run(cases, templates) + report = _report(runner) + assert report["counts"]["proposed"] == 1 + assert report["counts"]["rollout_attempted"] == 0 + assert ( + report["diagnostics"][report["audit"][0][0]["candidate_id"]][ + "planning_validation" + ]["checks"][0]["status"] + == "unavailable" + ) + assert runner.host.closed + + +@pytest.mark.parametrize("after", [False, True]) +def test_failure_counts_only_commands_actually_started_and_releases_capacity( + tmp_path, after +): + runner, cases, templates = _job(tmp_path) + runner.executor.fail_after, runner.executor.fail_before = after, not after + with pytest.raises(RuntimeError, match="command"): + runner.run(cases, templates) + report = _report(runner) + assert report["counts"]["rollout_attempted"] == int(after) + assert report["counts"]["committed"] == 0 + assert report["pending_reserved_bytes"] == 0 + assert runner.host.closed + + +@pytest.mark.parametrize("reason", ["task", "collision", "quality", "dynamics"]) +def test_actual_motion_gates_reject_before_expert_write(tmp_path, reason): + runner, cases, templates = _job( + tmp_path, + configuration={"collection": {"max_proposals": 1, "max_rollout_attempts": 1}}, + ) + if reason == "task": + runner.executor.status = "failed" + elif reason == "collision": + runner.planner.actual_status = "failed" + elif reason == "quality": + runner.executor.actual_middle = torch.tensor([-1.0, -1.0]) + else: + # Planned peak speed is 4; the measured excursion has speed 8. + runner.motion_limits.velocity_limits.fill_(5) + runner.executor.actual_middle = torch.tensor([0.4, 0.1]) + report = runner.run(cases, templates) + assert report["counts"]["rollout_attempted"] == 1 + assert report["counts"]["committed"] == 0 + assert report["stop_reason"] == "rollout_budget_exhausted" + assert not (runner.sink.root / "manifest.json").exists() + if reason == "collision": + assert len(runner.planner.calls) == 2 + assert torch.equal( + runner.planner.calls[-1].positions[0], + runner.executor.episodes[0].observations["joint_positions"], + ) + + +def test_real_persistence_failure_retries_without_another_rollout( + tmp_path, monkeypatch +): + runner, cases, templates = _job(tmp_path) + write = runner.sink._write_manifest + calls = [] + + def fail_once(record): + calls.append(record["commit_id"]) + if len(calls) == 1: + raise OSError("injected writer failure") + write(record) + + monkeypatch.setattr(runner.sink, "_write_manifest", fail_once) + report = runner.run(cases, templates) + assert calls[0] == calls[1] + assert report["counts"]["rollout_attempted"] == report["counts"]["committed"] == 1 + assert report["audit"][0][2] == 1 + assert report["pending_reserved_bytes"] == 0 + + +def test_sink_input_error_releases_pending_payload_and_preserves_error(tmp_path): + runner, cases, templates = _job(tmp_path) + runner.executor.unsupported_observation = True + with pytest.raises(ValueError, match="Unsupported observation"): + runner.run(cases, templates) + report = _report(runner) + assert report["counts"]["committed"] == 0 + assert report["pending_reserved_bytes"] == 0 + assert not (runner.sink.root / "manifest.json").exists() + + +def test_cancellation_stops_new_commands_and_commits_completed_row(tmp_path): + runner, cases, templates = _job( + tmp_path, rows=2, configuration={"collection": {"target_committed_episodes": 2}} + ) + stop = {"value": False} + runner.executor.after_first = lambda: stop.update(value=True) + report = runner.run(cases, templates, should_stop=lambda: stop["value"]) + assert report["cancelled"] + assert report["counts"]["rollout_attempted"] == report["counts"]["committed"] == 1 + assert report["pending_reserved_bytes"] == 0 + + +def test_restore_mismatch_stops_before_next_rollout(tmp_path): + runner, cases, templates = _job( + tmp_path, configuration={"collection": {"target_committed_episodes": 2}} + ) + runner.host.fail_restore = True + with pytest.raises(RuntimeError, match="restoration mismatch"): + runner.run(cases, templates) + report = _report(runner) + assert report["counts"]["rollout_attempted"] == report["counts"]["committed"] == 1 + assert report["pending_reserved_bytes"] == 0 + + +def test_template_validator_cannot_be_changed_by_job(tmp_path): + runner, cases, templates = _job(tmp_path) + with pytest.raises(ValueError, match="validator IDs"): + runner.run(cases, (replace(templates[0], validator_id="weaker"),)) + assert runner.host.prepared == 0 + assert runner.host.closed + + +@pytest.mark.parametrize("static_only", [False, True]) +def test_missing_rigid_collision_geometry_is_rejected_before_proposal( + tmp_path, static_only +): + runner, cases, templates = _job(tmp_path) + if static_only: + runner.planner.motion_generator.collision_world_entity_ids = ( + "unmodelled_cube", + ) + snapshots = runner.host.snapshots + runner.host.snapshots = lambda binding: tuple( + replace(value, entity_poses={"unmodelled_cube": torch.eye(4)}) + for value in snapshots(binding) + ) + with pytest.raises(ValueError, match="every rigid object"): + runner.run(cases, templates) + assert _report(runner)["counts"]["proposed"] == 0 + assert not runner.executor.calls + + +def test_wall_budget_expiring_during_planning_does_not_start_commands(tmp_path): + runner, cases, templates = _job(tmp_path) + clock = {"time": 0.0} + runner.clock = lambda: clock["time"] + validate = runner.planner.validate_qpos + + def slow_plan(*args): + clock["time"] = 61.0 + return validate(*args) + + runner.planner.validate_qpos = slow_plan + report = runner.run(cases, templates) + assert report["stop_reason"] == "wall_time_exhausted" + assert report["counts"]["rollout_attempted"] == 0 + assert report["pending_reserved_bytes"] == 0 + + +def test_report_failure_does_not_replace_original_execution_error( + tmp_path, monkeypatch +): + runner, cases, templates = _job(tmp_path) + runner.executor.fail_before = True + + def fail_report(report): + raise OSError("report disk failure") + + monkeypatch.setattr(runner, "_write_report", fail_report) + with pytest.raises(RuntimeError, match="before command") as caught: + runner.run(cases, templates) + assert "report disk failure" in caught.value.__notes__[0] + assert runner.host.closed + + +def test_cancel_during_second_planning_round_does_not_restore_or_execute(tmp_path): + runner, cases, templates = _job( + tmp_path, configuration={"collection": {"target_committed_episodes": 2}} + ) + stop = {"value": False} + validate = runner.planner.validate_qpos + + def cancel_on_next_plan(*args): + result = validate(*args) + if len(runner.planner.calls) == 3: # initial plan, actual path, second plan + stop["value"] = True + return result + + runner.planner.validate_qpos = cancel_on_next_plan + report = runner.run(cases, templates, should_stop=lambda: stop["value"]) + assert report["cancelled"] + assert runner.host.prepared == 1 + assert len(runner.executor.calls) == 1 + assert report["counts"]["committed"] == 1 + assert report["pending_reserved_bytes"] == 0 + + +def test_clock_mismatch_fails_before_initial_preparation(tmp_path): + runner, cases, templates = _job(tmp_path) + template = replace( + templates[0], + dt=torch.tensor([0.0, 0.05000005, 0.05000005], dtype=torch.float64), + ) + with pytest.raises(ValueError, match="timing must match"): + runner.run(cases, (template,)) + assert runner.host.prepared == 0 + assert not runner.executor.calls + + +def test_incomplete_transition_keeps_actual_backend_failure_in_audit(tmp_path): + runner, cases, templates = _job( + tmp_path, configuration={"collection": {"max_rollout_attempts": 1}} + ) + runner.executor.no_transition_error = "RuntimeError: physics integration failed" + report = runner.run(cases, templates) + assert report["counts"]["rollout_attempted"] == 1 + assert report["counts"]["committed"] == 0 + candidate_id = report["audit"][0][0].candidate_id + assert ( + report["diagnostics"][candidate_id]["reason"] + == runner.executor.no_transition_error + ) diff --git a/tests/lab/trajectory_generation/test_runner_real.py b/tests/lab/trajectory_generation/test_runner_real.py new file mode 100644 index 000000000..4a9cd7012 --- /dev/null +++ b/tests/lab/trajectory_generation/test_runner_real.py @@ -0,0 +1,101 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Measured UR5 acceptance and articulated Panda rejection through the real stack.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.requires_sim +@pytest.mark.parametrize("robot_type, accepted", [("ur5", True), ("panda", False)]) +def test_real_runner_accepts_only_validated_measured_free_motion( + tmp_path: Path, robot_type: str, accepted: bool +) -> None: + pytest.importorskip("curobo") + pytest.importorskip("lerobot") + av = pytest.importorskip("av") + from examples.sim.motion.trajectory_generation.free_motion import run_free_motion + + output = tmp_path / "collection" + report = run_free_motion(output, robot_type=robot_type, record_video=accepted) + assert report["counts"]["rollout_attempted"] == 1 + assert report["pending_reserved_bytes"] == 0 + assert not (output / ".writer.lock").exists() + if not accepted: + assert not report["target_reached"] + assert report["counts"]["committed"] == 0 + assert not (output / "manifest.json").exists() + assert not (output / "preview.mp4").exists() + diagnostics = next(iter(report["diagnostics"].values())) + checks = { + check.check_id: check for check in diagnostics["rollout_validation"].checks + } + assert checks["path_collision"].status == "unavailable" + assert checks["task_success"].metrics["max_gripper_drift"] > 1e-6 + return + assert report["target_reached"], report["diagnostics"] + assert report["counts"]["committed"] == report["counts"]["rollout_attempted"] == 1 + assert report["pending_reserved_bytes"] == 0 + assert not (output / ".writer.lock").exists() + manifest = json.loads((output / "manifest.json").read_text()) + assert len(manifest["episodes"]) == 1 + shard = output / manifest["episodes"][0]["shard"] + evidence = json.loads((shard / "episode.json").read_text()) + checks = {check["check_id"]: check for check in evidence["validation"]} + assert { + "planned_path_collision", + "path_collision", + "execution_complete", + "fixed_collision_world", + "task_success", + "actual_motion_limits", + "motion_quality", + } <= set(checks) + assert all(check["status"] == "passed" for check in checks.values()) + assert checks["task_success"]["metrics"]["measured_displacement"] >= 0.07 + assert checks["task_success"]["metrics"]["endpoint_error"] <= 0.01 + assert list((shard / "dataset").rglob("*.parquet")) + + # Rendering must retain every observation, including t=0 and the terminal + # state, without advancing the one-second physics rollout. + assert checks["motion_quality"]["metrics"]["duration_s"] == pytest.approx(1.0) + with np.load(shard / "terminal.npz") as terminal: + timestamps = terminal["timestamps"] + with av.open(str(output / "preview.mp4")) as video: + stream = video.streams.video[0] + assert stream.codec_context.name == "h264" + assert (stream.width, stream.height) == (640, 480) + assert stream.average_rate == evidence["fps"] == 20 + frames = list(video.decode(video=0)) + assert len(frames) == len(timestamps) == evidence["steps"] + 1 + np.testing.assert_allclose( + [float(frame.pts * frame.time_base) for frame in frames], + timestamps - timestamps[0], + rtol=0, + atol=1e-6, + ) + first = frames[0].to_ndarray(format="rgb24") + last = frames[-1].to_ndarray(format="rgb24") + assert first.std() > 1, "Camera preview must contain a rendered scene" + assert not np.array_equal(first, last), "Preview must show changing frames" diff --git a/tests/lab/trajectory_generation/test_session_sink_integration.py b/tests/lab/trajectory_generation/test_session_sink_integration.py new file mode 100644 index 000000000..e9a26f5a4 --- /dev/null +++ b/tests/lab/trajectory_generation/test_session_sink_integration.py @@ -0,0 +1,111 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Session accounting driven by actual sealed LeRobot persistence receipts.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CandidateTrajectoryBatch, + ExpertEpisode, + GenerationSession, + SceneCase, + TrajectoryGenerationJobCfg, + ValidationCheck, + ValidationResult, +) +from embodichain.lab.trajectory_generation.sinks import LeRobotEpisodeSink + + +def test_failed_real_commit_retry_preserves_rollout_count_and_confirms_coverage( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # Synthetic rollout evidence isolates receipt/accounting from physical quality. + cfg = TrajectoryGenerationJobCfg.from_mapping( + {"collection": {"target_committed_episodes": 1}} + ) + session = GenerationSession(cfg) + case = SceneCase("case", "initial", "scene", "move", "robot") + session.register_case( + case, torch.tensor([[-1.0, 1.0], [-1.0, 1.0]]), joint_names=("a", "b") + ) + identity, _ = session.propose( + "case", + "initial", + source_id="test_source", + source_revision="v1", + template_id="reference", + operator_id="replay", + ) + positions = torch.tensor([[[0.0, 0.0], [0.25, 0.0], [0.5, 0.0]]]) + dt = torch.tensor([[0.0, 0.05, 0.05]], dtype=torch.float64) + plan = CandidateTrajectoryBatch( + positions, dt, torch.tensor([3]), (identity,), ("a", "b") + ) + session.add_planned( + plan, ValidationResult((ValidationCheck("path_collision", "passed"),)) + ) + assert session.take_ready("case", "initial", episode_byte_budget=4096) is not None + session.mark_rollout_started(identity) + episode_id, commit_id = session.episode_ids(identity) + episode = ExpertEpisode( + identity, + {"joint_positions": positions[0]}, + positions[0, 1:], + dt[0].cumsum(0), + "qpos", + ValidationResult( + ( + ValidationCheck("path_collision", "passed"), + ValidationCheck("task_success", "passed"), + ) + ), + episode_id, + commit_id, + ) + assert session.accept_episode(episode) + with LeRobotEpisodeSink(tmp_path / "episodes", fps=20) as sink: + write_manifest = sink._write_manifest + + def fail_manifest(record): + raise OSError("manifest not written") + + monkeypatch.setattr(sink, "_write_manifest", fail_manifest) + failed = sink.submit(episode) + assert not failed.confirmed + session.apply_receipt(failed) + assert session.snapshot()["counts"]["committed"] == 0 + assert session.snapshot()["coverage"]["case"] == 0 + retry_episode, submission = session.retry_write(commit_id) + monkeypatch.setattr(sink, "_write_manifest", write_manifest) + confirmed = sink.submit(retry_episode, submission_id=submission) + assert confirmed.confirmed, confirmed.error + session.apply_receipt(confirmed) + session.apply_receipt(confirmed) + report = session.snapshot() + assert report["counts"]["committed"] == 1 + assert report["counts"]["rollout_attempted"] == 1 + assert report["coverage"]["case"] == 1 + assert report["pending_reserved_bytes"] == 0 + manifest = json.loads((sink.root / "manifest.json").read_text()) + assert len(manifest["episodes"]) == 1 + assert manifest["episodes"][0]["commit_id"] == episode.commit_id diff --git a/tests/lab/trajectory_generation/test_sim_initial_state.py b/tests/lab/trajectory_generation/test_sim_initial_state.py new file mode 100644 index 000000000..14d5d0322 --- /dev/null +++ b/tests/lab/trajectory_generation/test_sim_initial_state.py @@ -0,0 +1,463 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Physical initial-state contracts without creating a simulation world.""" + +from __future__ import annotations + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.lab.trajectory_generation.integrations.sim import ( + SimInitialStateAdapter, +) + +_BATCH = 2 +_JOINTS = 3 # Include one mimic/gripper coordinate in addition to the arm. + + +class _Robot: + """Mutable full-joint state with the production Robot getter/setter API.""" + + def __init__(self) -> None: + self.uid = "robot" + self.device = torch.device("cpu") + self.num_instances = _BATCH + self.dof = _JOINTS + self.joint_names = ["arm", "gripper", "mimic"] + self.control_parts = {"arm": ["arm"], "gripper": ["gripper", "mimic"]} + self.cfg = SimpleNamespace(fix_base=True, use_usd_properties=False) + self.pose = torch.eye(4).repeat(_BATCH, 1, 1) + self.pose[1, 0, 3] = 2.0 + self.qpos = torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) + self.qvel = self.qpos / 10 + self.target_qpos = self.qpos + 0.01 + self.target_qvel = self.qpos / 20 + self.qf = self.qpos / 30 + self.limits = torch.tensor([-1.0, 1.0]).repeat(_BATCH, _JOINTS, 1) + self.root_link_name = "base" + velocities = torch.zeros(_BATCH, 2, 6) + self.body_data = SimpleNamespace( + link_names=["tip", "base"], # Root selection must use identity. + entities=[SimpleNamespace(get_root_link_name=lambda: "base")], + body_link_vel=velocities, + root_lin_vel=velocities[:, 1, :3], + root_ang_vel=velocities[:, 1, 3:], + ) + self.writes: list[str] = [] + self.after_root_write = lambda: None + + def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: + assert to_matrix + return self.pose + + def get_qpos(self, target: bool = False) -> torch.Tensor: + return self.target_qpos if target else self.qpos + + def get_qvel(self, target: bool = False) -> torch.Tensor: + return self.target_qvel if target else self.qvel + + def get_qf(self) -> torch.Tensor: + return self.qf + + def get_qpos_limits(self) -> torch.Tensor: + return self.limits + + def set_local_pose(self, value: torch.Tensor) -> None: + self.writes.append("root_pose") + self.pose.copy_(value) + self.after_root_write() + + def set_qpos(self, value: torch.Tensor, target: bool = True) -> None: + self.writes.append("target_qpos" if target else "qpos") + self.get_qpos(target=target).copy_(value) + + def set_qvel(self, value: torch.Tensor, target: bool = True) -> None: + self.writes.append("target_qvel" if target else "qvel") + self.get_qvel(target=target).copy_(value) + + def set_qf(self, value: torch.Tensor) -> None: + self.writes.append("qf") + self.qf.copy_(value) + + +class _RigidObject: + """Minimal rigid body retaining force state independently of velocity.""" + + def __init__(self, uid: str, *, static: bool = False) -> None: + self.uid = uid + self.device = torch.device("cpu") + self.num_instances = _BATCH + self.is_static = self.is_non_dynamic = static + self.pose = torch.eye(4).repeat(_BATCH, 1, 1) + self.pose[:, 0, 3] = torch.tensor([0.2, 0.8]) + self.body_state = torch.zeros(_BATCH, 13) + if not static: + self.body_state[:, 7:] = torch.arange(12).reshape(_BATCH, 6) / 100 + self.pending_force = torch.zeros(_BATCH, 3) + self.writes: list[str] = [] + + def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: + assert to_matrix + return self.pose + + def set_local_pose(self, value: torch.Tensor) -> None: + self.writes.append("pose") + self.pose.copy_(value) + + def clear_dynamics(self) -> None: + self.writes.append("clear_dynamics") + self.body_state[:, 7:] = 0 + self.pending_force.zero_() + + def set_velocity(self, lin_vel: torch.Tensor, ang_vel: torch.Tensor) -> None: + self.writes.append("velocity") + self.body_state[:, 7:10] = lin_vel + self.body_state[:, 10:13] = ang_vel + + +def _scene() -> tuple[SimpleNamespace, _Robot, SimInitialStateAdapter]: + robot = _Robot() + sim = SimpleNamespace( + num_envs=_BATCH, + _robots={robot.uid: robot}, + _rigid_objects={ + "cube": _RigidObject("cube"), + "table": _RigidObject("table", static=True), + }, + _articulations={}, + _rigid_object_groups={}, + _soft_objects={}, + _cloth_objects={}, + _constraints={}, + ) + return sim, robot, SimInitialStateAdapter(sim, robot) + + +def test_capture_owns_nested_tensor_state_and_complete_joint_order() -> None: + sim, robot, adapter = _scene() + state = adapter.capture() + robot.qpos.zero_() + sim._rigid_objects["cube"].body_state.zero_() + assert state.joint_names == ("arm", "gripper", "mimic") + assert state.robot["qpos"][1, 2] == pytest.approx(0.6) + assert state.rigid_objects["cube"]["angular_velocity"][1, 2] == pytest.approx(0.11) + with pytest.raises(TypeError): + state.robot["qpos"] = torch.zeros(_BATCH, _JOINTS) + with pytest.raises(TypeError): + state.rigid_objects["unknown"] = {} + + +def test_restore_repeatedly_recovers_all_rows_targets_and_pending_forces() -> None: + sim, robot, adapter = _scene() + state = adapter.capture() + cube = sim._rigid_objects["cube"] + + def advance_world_during_root_write() -> None: + # The real root setter advances the world before remaining restoration. + robot.qpos.add_(0.1) + cube.pose[:, 0, 3] += 0.3 + cube.body_state[:, 7:] += 0.2 + + robot.after_root_write = advance_world_during_root_write + for _ in range(3): + robot.qpos.zero_() + robot.target_qpos.zero_() + robot.qvel.zero_() + robot.target_qvel.zero_() + robot.qf.zero_() + cube.pending_force.fill_(4) + adapter.restore(state) + assert adapter.verify(state).accepted + assert not cube.pending_force.any() + assert robot.writes[-6:] == [ + "root_pose", + "qpos", + "qvel", + "target_qpos", + "target_qvel", + "qf", + ] + assert sim._rigid_objects["table"].writes[-1:] == ["pose"] + + +@pytest.mark.parametrize("field", ["qpos", "qvel", "target_qpos", "target_qvel", "qf"]) +def test_verify_detects_joint_or_controller_drift(field: str) -> None: + _, robot, adapter = _scene() + state = adapter.capture() + getattr(robot, field)[1, 2] += 0.1 + result = adapter.verify(state) + assert not result.accepted + assert f"robot.{field}" in result.checks[0].detail + + +def test_verify_detects_root_velocity_and_rigid_velocity_drift() -> None: + sim, robot, adapter = _scene() + state = adapter.capture() + robot.body_data.root_lin_vel[0, 0] = 0.1 + assert not adapter.verify(state).accepted + robot.body_data.root_lin_vel.zero_() + sim._rigid_objects["cube"].body_state[1, 10] += 0.1 + assert not adapter.verify(state).accepted + + +def test_verify_nonfinite_live_state_fails_without_physical_write() -> None: + _, robot, adapter = _scene() + state = adapter.capture() + robot.qvel[0, 0] = float("nan") + assert not adapter.verify(state).accepted + assert robot.writes == [] + + +@pytest.mark.parametrize( + "registry", + [ + "_articulations", + "_rigid_object_groups", + "_soft_objects", + "_cloth_objects", + "_constraints", + ], +) +def test_unsupported_entity_or_constraint_rejected_before_restore( + registry: str, +) -> None: + sim, robot, adapter = _scene() + state = adapter.capture() + getattr(sim, registry)["unsupported"] = object() + with pytest.raises(ValueError, match="does not support"): + adapter.restore(state) + assert robot.writes == [] + + +@pytest.mark.parametrize( + "change", + [ + "extra_robot", + "partial_batch", + "floating_base", + "usd_base", + "joint_order", + "object_uid", + "control_parts", + "body_mode", + ], +) +def test_topology_and_control_changes_fail_preflight(change: str) -> None: + sim, robot, adapter = _scene() + state = adapter.capture() + if change == "extra_robot": + sim._robots["other"] = _Robot() + elif change == "partial_batch": + robot.num_instances = 1 + elif change == "floating_base": + robot.cfg.fix_base = False + elif change == "usd_base": + robot.cfg.use_usd_properties = True + elif change == "joint_order": + robot.joint_names.reverse() + elif change == "object_uid": + sim._rigid_objects["cube"].uid = "renamed" + elif change == "control_parts": + robot.control_parts["arm"] = ["gripper"] + else: + sim._rigid_objects["cube"].is_non_dynamic = True + with pytest.raises(ValueError): + adapter.restore(state) + assert robot.writes == [] + + +@pytest.mark.parametrize( + "change", + [ + "missing_field", + "wrong_shape", + "nonfinite", + "nonrigid_pose", + "root_velocity", + "joint_limit", + "object_missing", + ], +) +def test_malformed_snapshot_rejected_before_any_write(change: str) -> None: + _, robot, adapter = _scene() + state = adapter.capture() + if change == "missing_field": + state = replace( + state, + robot={ + key: value for key, value in state.robot.items() if key != "target_qvel" + }, + ) + elif change == "wrong_shape": + state = replace(state, robot={**state.robot, "qpos": torch.zeros(1, _JOINTS)}) + elif change == "nonfinite": + state.robot["qpos"][0, 0] = float("nan") + elif change == "nonrigid_pose": + state.robot["root_pose"][0, 0, 0] = -1 + elif change == "root_velocity": + state.robot["root_linear_velocity"][0, 0] = 0.1 + elif change == "joint_limit": + state.robot["target_qpos"][0, 0] = 2 + else: + state = replace(state, rigid_objects={}) + with pytest.raises(ValueError): + adapter.restore(state) + assert robot.writes == [] + + +def test_signature_tracks_structure_and_ignores_episode_motion() -> None: + sim, robot, adapter = _scene() + signature = adapter.signature() + robot.qpos.add_(0.1) + sim._rigid_objects["cube"].pose[:, 0, 3] += 1 + assert adapter.signature() == signature + sim._rigid_objects = dict(reversed(tuple(sim._rigid_objects.items()))) + assert adapter.signature() == signature + + +def test_verification_uses_explicit_absolute_tolerance() -> None: + sim, robot, _ = _scene() + tolerance = 0.01 + adapter = SimInitialStateAdapter(sim, robot, atol=tolerance) + state = adapter.capture() + robot.qpos[0, 0] += tolerance / 2 + assert adapter.verify(state).accepted + robot.qpos[0, 0] += tolerance + assert not adapter.verify(state).accepted + + +@pytest.mark.parametrize("atol", [-1, float("nan"), float("inf"), True]) +def test_invalid_tolerance_rejected(atol: float) -> None: + sim, robot, _ = _scene() + with pytest.raises((ValueError, TypeError)): + SimInitialStateAdapter(sim, robot, atol=atol) + + +@pytest.mark.requires_sim +@pytest.mark.parametrize("physics_enabled", [False, True]) +def test_real_cpu_robot_and_rigid_object_initial_state( + tmp_path, physics_enabled +) -> None: + """Exercise the actual CPU setters/getters without assets, IK, or cameras.""" + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.cfg import RigidObjectCfg, RobotCfg + from embodichain.lab.sim.shapes import CubeCfg + + urdf = tmp_path / "initial_state_robot.urdf" + urdf.write_text( + '' + '' + '' + "" + '' + '' + '' + '' + '' + '' + '', + encoding="utf-8", + ) + sim = SimulationManager( + SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=1) + ) + try: + sim.enable_physics(physics_enabled) + robot = sim.add_robot( + RobotCfg(uid="initial_robot", fpath=str(urdf), fix_base=True) + ) + cube = sim.add_rigid_object( + RigidObjectCfg( + uid="initial_cube", + shape=CubeCfg(size=[0.05, 0.05, 0.05]), + init_pos=[1, 0, 1], + ) + ) + # Distinct nonzero values catch confusion between current and target APIs. + robot.set_qpos(torch.tensor([[0.2]]), target=False) + robot.set_qpos(torch.tensor([[0.3]]), target=True) + robot.set_qvel(torch.tensor([[0.1]]), target=False) + robot.set_qvel(torch.tensor([[0.15]]), target=True) + robot.set_qf(torch.tensor([[0.05]])) + cube.set_velocity( + lin_vel=torch.tensor([[0.1, 0.2, 0.3]]), + ang_vel=torch.tensor([[0.3, 0.2, 0.1]]), + ) + adapter = SimInitialStateAdapter(sim, robot) + state = adapter.capture() + robot.set_qpos(torch.tensor([[-0.2]]), target=False) + robot.set_qpos(torch.tensor([[-0.3]]), target=True) + robot.set_qvel(torch.zeros(1, 1), target=False) + robot.set_qvel(torch.zeros(1, 1), target=True) + robot.set_qf(torch.zeros(1, 1)) + cube.clear_dynamics() + pose = cube.get_local_pose(to_matrix=True) + pose[:, 0, 3] += 0.2 + cube.set_local_pose(pose) + adapter.restore(state) + result = adapter.verify(state) + assert result.accepted, result.checks + finally: + sim.destroy() + SimulationManager.flush_cleanup_queue() + + +def test_host_planning_snapshots_own_state_and_keep_physical_case_rows() -> None: + from embodichain.lab.sim.motion.trajectory_augmentation import ( + SceneCase, + ValidationCheck, + ValidationResult, + ) + from embodichain.lab.trajectory_generation.initial_state import ( + FixedSceneHost, + InitialStateProfile, + ) + + sim, robot, adapter = _scene() + cases = tuple( + SceneCase(f"case-{row}", f"initial-{row}", "scene", "move", "robot") + for row in range(_BATCH) + ) + profile = InitialStateProfile( + profile_id="fixed_scene_initial_state", + prepare=lambda: None, + signature=lambda: "fixed-configuration", + verify=lambda cases: ValidationResult( + (ValidationCheck("task_initial", "passed"),) + ), + ) + with FixedSceneHost(adapter, profile) as host: + binding = host.acquire_case(cases) + snapshots = host.snapshots(binding) + assert [value.scene_case for value in snapshots] == list(cases) + assert torch.equal(snapshots[1].joint_positions, robot.qpos[1]) + assert torch.equal(snapshots[1].root_pose, robot.pose[1]) + expected = snapshots[0].joint_positions.clone() + snapshots[0].joint_positions.zero_() + snapshots[0].entity_poses["cube"].zero_() + fresh = host.snapshots(binding) + assert torch.equal(fresh[0].joint_positions, expected) + assert torch.equal( + fresh[0].entity_poses["cube"], sim._rigid_objects["cube"].pose[0] + ) + assert host.initial_observation(binding) is None + host.restore_initial() + with pytest.raises(RuntimeError, match="obsolete"): + host.snapshots(binding) diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 11b46bace..b34a87dea 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -108,7 +108,7 @@ ParallelJawGraspPoseGenerator, ParallelJawGripperModelCfg, ) -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( MotionGenerator, MoveType, PlanOptions, @@ -263,7 +263,7 @@ def interpolate( interpolate, ) monkeypatch.setattr( - "embodichain.lab.sim.planners.motion_generator.interpolate_with_distance", + "embodichain.lab.sim.motion.planners.motion_generator.interpolate_with_distance", interpolate, ) monkeypatch.setattr( @@ -2068,6 +2068,252 @@ def compute_batch_ik( assert qpos.shape == (NUM_ENVS, 3, 2, ARM_DOF) +@pytest.mark.parametrize("all_empty", (False, True)) +def test_pick_empty_grasp_rows_return_safe_pose_and_failure_mask( + all_empty: bool, +) -> None: + generator = _motion_generator() + action = _bind_action(generator, PickUp()) + empty = (torch.empty(0, 4, 4), torch.empty(0)) + valid_pose = torch.eye(4).unsqueeze(0) + valid_pose[:, 0, 3] = 0.3 + _GRASP_GENERATORS[id(action)].get_valid_grasp_poses = Mock( + return_value=[empty, empty if all_empty else (valid_pose, torch.zeros(1))] + ) + action._select_feasible_grasp_variants = Mock( + side_effect=lambda poses, *_: ( + poses, + torch.ones(poses.shape[:2], dtype=torch.bool), + ) + ) + safe_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + safe_pose[:, 1, 3] = 0.4 + generator.robot.compute_fk.side_effect = None + generator.robot.compute_fk.return_value = safe_pose + + success, selected = action._resolve_grasp_pose( + ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, entity_id="target" + ), + torch.eye(4).repeat(NUM_ENVS, 1, 1), + torch.zeros(NUM_ENVS, ARM_DOF), + JointPositionTarget("arm", tuple(range(ARM_DOF))), + "hand", + PickUpOptions(), + torch.tensor((0.0, 0.0, -1.0)), + ) + + assert success.tolist() == [False, not all_empty] + torch.testing.assert_close(selected[0], safe_pose[0]) + if all_empty: + action._select_feasible_grasp_variants.assert_not_called() + torch.testing.assert_close(selected, safe_pose) + else: + passed_poses = action._select_feasible_grasp_variants.call_args.args[0] + torch.testing.assert_close(passed_poses[0, 0], safe_pose[0]) + torch.testing.assert_close(selected[1], valid_pose[0]) + + +def test_pick_padding_cannot_become_a_feasible_grasp() -> None: + generator = _motion_generator() + action = _bind_action(generator, PickUp()) + _GRASP_GENERATORS[id(action)].get_valid_grasp_poses = Mock( + return_value=[ + (torch.eye(4).repeat(1, 1, 1), torch.zeros(1)), + (torch.eye(4).repeat(2, 1, 1), torch.zeros(2)), + ] + ) + action._select_feasible_grasp_variants = Mock( + side_effect=lambda poses, *_: ( + poses, + torch.tensor([[False, True], [True, True]]), + ) + ) + success, _ = action._resolve_grasp_pose( + ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, entity_id="target" + ), + torch.eye(4).repeat(NUM_ENVS, 1, 1), + torch.zeros(NUM_ENVS, ARM_DOF), + JointPositionTarget("arm", tuple(range(ARM_DOF))), + "hand", + PickUpOptions(), + torch.tensor((0.0, 0.0, -1.0)), + ) + assert success.tolist() == [False, True] + + +@pytest.mark.parametrize( + "invalid_kind", + ( + "nan_cost", + "infinite_cost", + "negative_infinite_cost", + "nan_pose", + "infinite_pose", + "reflection", + "nonorthogonal_rotation", + "invalid_bottom_row", + ), +) +def test_pick_invalid_grasp_never_reaches_ik_or_hides_valid_peer( + invalid_kind: str, +) -> None: + generator = _motion_generator() + action = _bind_action(generator, PickUp()) + invalid_pose = torch.eye(4) + invalid_pose[0, 3] = 0.1 + invalid_cost = 0.0 + if invalid_kind == "nan_cost": + invalid_cost = float("nan") + elif invalid_kind == "infinite_cost": + invalid_cost = float("inf") + elif invalid_kind == "negative_infinite_cost": + invalid_cost = -float("inf") + elif invalid_kind == "nan_pose": + invalid_pose[0, 3] = float("nan") + elif invalid_kind == "infinite_pose": + invalid_pose[0, 0] = float("inf") + elif invalid_kind == "reflection": + invalid_pose[0, 0] = -1.0 + elif invalid_kind == "nonorthogonal_rotation": + invalid_pose[0, 0] = 2.0 + else: + invalid_pose[3, 0] = 1.0 + valid_pose = torch.eye(4) + valid_pose[0, 3] = 0.3 + _GRASP_GENERATORS[id(action)].get_valid_grasp_poses = Mock( + return_value=[ + ( + torch.stack((invalid_pose, valid_pose)), + torch.tensor([invalid_cost, 1.0]), + ), + (invalid_pose.unsqueeze(0), torch.tensor([invalid_cost])), + ] + ) + safe_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + safe_pose[:, 1, 3] = 0.4 + generator.robot.compute_fk.side_effect = None + generator.robot.compute_fk.return_value = safe_pose + + def feasible_candidates( + poses: torch.Tensor, *_: object + ) -> tuple[torch.Tensor, torch.Tensor]: + assert torch.isfinite(poses).all() + torch.testing.assert_close(poses[:, 0], safe_pose) + return poses, torch.ones(poses.shape[:2], dtype=torch.bool) + + action._select_feasible_grasp_variants = Mock(side_effect=feasible_candidates) + success, selected = action._resolve_grasp_pose( + ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, entity_id="target" + ), + torch.eye(4).repeat(NUM_ENVS, 1, 1), + torch.zeros(NUM_ENVS, ARM_DOF), + JointPositionTarget("arm", tuple(range(ARM_DOF))), + "hand", + PickUpOptions(), + torch.tensor((0.0, 0.0, -1.0)), + ) + + assert success.tolist() == [True, False] + torch.testing.assert_close(selected[0], valid_pose) + torch.testing.assert_close(selected[1], safe_pose[1]) + + +@pytest.mark.parametrize("cost_offset", (0.0, 10000.0)) +def test_pick_finite_costs_choose_stable_minimum_without_magic_threshold( + cost_offset: float, +) -> None: + generator = _motion_generator() + action = _bind_action(generator, PickUp()) + poses = torch.eye(4).repeat(3, 1, 1) + poses[:, 0, 3] = torch.tensor([0.1, 0.2, 0.3]) + _GRASP_GENERATORS[id(action)].get_valid_grasp_poses = Mock( + return_value=[ + (poses, torch.tensor([2.0, 1.0, 1.0]) + cost_offset), + (poses, torch.tensor([0.5, 2.0, 1.0]) + cost_offset), + ] + ) + action._select_feasible_grasp_variants = Mock( + side_effect=lambda poses, *_: ( + poses, + torch.ones(poses.shape[:2], dtype=torch.bool), + ) + ) + success, selected = action._resolve_grasp_pose( + ObjectSemantics( + affordance=AntipodalAffordance(), geometry={}, entity_id="target" + ), + torch.eye(4).repeat(NUM_ENVS, 1, 1), + torch.zeros(NUM_ENVS, ARM_DOF), + JointPositionTarget("arm", tuple(range(ARM_DOF))), + "hand", + PickUpOptions(), + torch.tensor((0.0, 0.0, -1.0)), + ) + assert success.all() + torch.testing.assert_close(selected, poses[torch.tensor([1, 0])]) + + +@pytest.mark.parametrize("bad_qpos", (float("nan"), float("inf"), 99.0)) +def test_pick_failed_ik_preserves_last_valid_seed(bad_qpos: float) -> None: + generator = _motion_generator() + action = _bind_action(generator, PickUp()) + poses = torch.eye(4).repeat(NUM_ENVS, 1, 2, 1, 1) + seed = torch.full((NUM_ENVS, ARM_DOF), 0.25) + + def compute_batch_ik( + *, pose: torch.Tensor, name: str, joint_seed: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + success = torch.tensor([[False, True], [True, True]]) + qpos = joint_seed + 0.1 + qpos[0, 0] = bad_qpos + # A backend incorrectly flagging a NaN result as success must also fail. + qpos[1, 0] = float("nan") + return success, qpos + + generator.robot.compute_batch_ik.side_effect = compute_batch_ik + success, qpos = action._compute_batch_candidate_ik( + poses, seed, JointPositionTarget("arm", tuple(range(ARM_DOF))) + ) + assert success.tolist() == [[[False, True]], [[False, True]]] + torch.testing.assert_close(qpos[:, 0, 0], seed) + torch.testing.assert_close(qpos[:, 0, 1], seed + 0.1) + assert torch.isfinite(qpos).all() + + +def test_pick_failed_pregrasp_ik_never_passes_nan_seed_to_later_stages() -> None: + generator = _motion_generator() + action = _bind_action(generator, PickUp()) + seeds: list[torch.Tensor] = [] + + def compute_batch_ik( + *, pose: torch.Tensor, name: str, joint_seed: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + assert torch.isfinite(joint_seed).all() + seeds.append(joint_seed.clone()) + success = torch.ones(pose.shape[:2], dtype=torch.bool) + qpos = joint_seed + 0.1 + if len(seeds) == 1: + success[0] = False + qpos[0] = float("nan") + return success, qpos + + generator.robot.compute_batch_ik.side_effect = compute_batch_ik + _, success = action._select_feasible_grasp_variants( + torch.eye(4).repeat(NUM_ENVS, 1, 1, 1), + torch.zeros(NUM_ENVS, ARM_DOF), + torch.eye(4).repeat(NUM_ENVS, 1, 1), + JointPositionTarget("arm", tuple(range(ARM_DOF))), + PickUpOptions(downstream_object_target_poses=(torch.eye(4),)), + torch.tensor((0.0, 0.0, -1.0)), + ) + assert success.tolist() == [[False], [True]] + assert len(seeds) == 4 # Pregrasp, grasp, lift, and downstream target. + torch.testing.assert_close(seeds[1][0], torch.zeros_like(seeds[1][0])) + + def test_pick_maps_canonical_grasp_frames_to_the_robot_eef() -> None: """Endpoint calibration is applied after canonical grasp generation.""" generator = _motion_generator() diff --git a/tests/sim/atomic_actions/test_core.py b/tests/sim/atomic_actions/test_core.py index e4b2ffa90..6c5ad2352 100644 --- a/tests/sim/atomic_actions/test_core.py +++ b/tests/sim/atomic_actions/test_core.py @@ -74,7 +74,7 @@ collect_scene_dependencies, resolve_pose_goal, ) -from embodichain.lab.sim.planners import ToppraPlanOptions +from embodichain.lab.sim.motion.planners import ToppraPlanOptions def _semantics( diff --git a/tests/sim/atomic_actions/test_curobo_motion_strategy_e2e.py b/tests/sim/atomic_actions/test_curobo_motion_strategy_e2e.py index 7ea946be9..9411f05d8 100644 --- a/tests/sim/atomic_actions/test_curobo_motion_strategy_e2e.py +++ b/tests/sim/atomic_actions/test_curobo_motion_strategy_e2e.py @@ -37,11 +37,11 @@ from embodichain.lab.sim.objects import RigidObjectCfg # noqa: E402 from embodichain.lab.sim.robots import FrankaPandaCfg # noqa: E402 from embodichain.lab.sim.shapes import CubeCfg # noqa: E402 -from embodichain.lab.sim.planners import ( # noqa: E402 +from embodichain.lab.sim.motion.planners import ( # noqa: E402 MotionGenCfg, MotionGenerator, ) -from embodichain.lab.sim.planners.curobo.curobo_planner import ( # noqa: E402 +from embodichain.lab.sim.motion.planners.curobo.curobo_planner import ( # noqa: E402 CuroboPlannerCfg, CuroboWorldCfg, ) diff --git a/tests/sim/atomic_actions/test_endpoint_runtime_e2e.py b/tests/sim/atomic_actions/test_endpoint_runtime_e2e.py index 9687ee143..7ffec130b 100644 --- a/tests/sim/atomic_actions/test_endpoint_runtime_e2e.py +++ b/tests/sim/atomic_actions/test_endpoint_runtime_e2e.py @@ -55,7 +55,7 @@ TrackingPolicy, ) from embodichain.lab.sim.atomic_actions.invocation import ResolvedActionRequest -from embodichain.lab.sim.planners import PlanResult +from embodichain.lab.sim.motion.planners import PlanResult from embodichain.lab.task_program.semantics import ( EndpointResolution, ResourceBinding, diff --git a/tests/sim/atomic_actions/test_engine_per_env.py b/tests/sim/atomic_actions/test_engine_per_env.py index 5df2b2dbf..7f5cb6041 100644 --- a/tests/sim/atomic_actions/test_engine_per_env.py +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -84,7 +84,7 @@ TrajectorySegment, ) from embodichain.lab.sim.atomic_actions.goals import resolve_pose_goal -from embodichain.lab.sim.planners import PlanOptions +from embodichain.lab.sim.motion.planners import PlanOptions CONTROL_DT = 0.02 @@ -932,6 +932,237 @@ def test_session_completes_incremental_command_sequence() -> None: assert final.eligible_mask.tolist() == [True] +def _candidate_plan( + action: DynamicAction, + request: ResolvedActionRequest, + context: PlanningContext, +) -> ActionPlan: + """Materialize a selected candidate without invoking the skill planner.""" + target = torch.full_like(context.robot.qpos, 0.7) + return action.build_plan( + request, + context, + success=True, + trajectory=TimedTrajectory.from_uniform_step( + torch.stack((context.robot.qpos, target), dim=1), + env_ids=context.env_ids, + step_dt=0.1, + ), + ) + + +def test_initial_plan_provider_consumes_candidate_then_plans_later_invocation() -> None: + engine, action = _engine() + initial = _context(0.0, 0.0, 0.2, 0) + invocation = _invocation(engine) + provider = Mock( + side_effect=lambda request, context: _candidate_plan(action, request, context) + ) + + session = engine.start( + (invocation, replace(invocation, invocation_id="next-call")), + initial, + initial_plan_provider=provider, + ) + assert action.plan_count == 0 + request, measured = provider.call_args.args + assert isinstance(request, ResolvedActionRequest) + assert measured is initial + session.tick(initial) + selected = session.tick(_context(0.1, 0.0, 0.2, 0)) + torch.testing.assert_close( + _joint_positions(selected.command), torch.full((1, 2), 0.7) + ) + session.tick(_context(0.2, 0.7, 0.2, 0)) + restarted = session.tick(_context(0.3, 0.7, 0.2, 0)) + torch.testing.assert_close( + _joint_positions(restarted.command), torch.full((1, 2), 0.7) + ) + following = session.tick(_context(0.4, 0.7, 0.2, 0)) + torch.testing.assert_close( + _joint_positions(following.command), torch.full((1, 2), 0.2) + ) + completed = session.tick(_context(0.5, 0.2, 0.2, 0)) + assert completed.status is ExecutionStatus.COMPLETED + assert action.plan_count == 1 + assert action.requests[0].invocation_id == "next-call" + provider.assert_called_once() + + +def test_initial_plan_provider_does_not_replace_recovery_planner() -> None: + engine, action = _engine() + initial = _context(0.0, 0.0, 0.2, 0) + provider = Mock( + side_effect=lambda request, context: _candidate_plan(action, request, context) + ) + session = engine.start( + (_invocation(engine),), initial, initial_plan_provider=provider + ) + session.tick(initial) + + recovered = session.tick(_context(0.1, 0.0, 0.4, 1)) + assert ExecutionEventKind.REPLANNED in {event.kind for event in recovered.events} + assert action.plan_count == 1 + assert session.plan_attempts[-1].plan.planned_scene_version == 1 + provider.assert_called_once() + + +@pytest.mark.parametrize( + ("field", "value", "message"), + ( + ("skill_id", "another-skill", "skill_id must match"), + ("invocation_id", "old-call", "correlation id"), + ("invocation_revision", 1, "request revision"), + ("planned_scene_version", 1, "planning scene version"), + ("planned_collision_world_revision", (1,), "collision-world revision"), + ), +) +def test_initial_plan_provider_rejects_mismatched_plan( + field: str, value: object, message: str +) -> None: + engine, action = _engine() + + def provider( + request: ResolvedActionRequest, context: PlanningContext + ) -> ActionPlan: + return replace(_candidate_plan(action, request, context), **{field: value}) + + with pytest.raises(ValueError, match=message): + engine.start( + (_invocation(engine),), + _context(0.0, 0.0, 0.2, 0), + initial_plan_provider=provider, + ) + assert action.plan_count == 0 + + +def test_initial_plan_provider_rejects_wrong_result_type() -> None: + engine, _ = _engine() + with pytest.raises(TypeError, match="must return an ActionPlan"): + engine.start( + (_invocation(engine),), + _context(0.0, 0.0, 0.2, 0), + initial_plan_provider=Mock(return_value=None), + ) + + +def test_initial_plan_provider_cannot_command_an_unbound_target() -> None: + engine, action = _engine() + + def provider( + request: ResolvedActionRequest, context: PlanningContext + ) -> ActionPlan: + plan = _candidate_plan(action, request, context) + frames = tuple( + replace( + frame, + commands=tuple( + replace(command, target=JointPositionTarget("unbound", (0, 1))) + for command in frame.commands + ), + ) + for frame in plan.commands.frames + ) + return replace(plan, commands=replace(plan.commands, frames=frames)) + + with pytest.raises(ValueError, match="not authorized"): + engine.start( + (_invocation(engine),), + _context(0.0, 0.0, 0.2, 0), + initial_plan_provider=provider, + ) + + +def test_initial_plan_provider_receives_new_request_after_reset() -> None: + engine, action = _engine() + invocation = _invocation(engine) + provider = Mock( + side_effect=lambda request, context: _candidate_plan(action, request, context) + ) + before = _context(1.0, 0.3, 0.2, 5) + after = _context(0.0, 0.0, 0.2, 0) + + first = engine.start((invocation,), before, initial_plan_provider=provider) + second = engine.start((invocation,), after, initial_plan_provider=provider) + + first_request, first_context = provider.call_args_list[0].args + second_request, second_context = provider.call_args_list[1].args + assert first_request is not second_request + assert first_context is before + assert second_context is after + torch.testing.assert_close( + _joint_positions(first.tick(before).command), before.robot.qpos + ) + torch.testing.assert_close( + _joint_positions(second.tick(after).command), after.robot.qpos + ) + assert action.plan_count == 0 + + +def test_initial_plan_provider_preserves_required_collision_binding() -> None: + engine, action = _engine() + provider = Mock( + side_effect=lambda request, context: _candidate_plan(action, request, context) + ) + with pytest.raises(ValueError, match="dynamic_collision_mode='required'"): + engine.start( + ( + _invocation( + engine, dynamic_collision_mode=DynamicCollisionMode.REQUIRED + ), + ), + _context(0.0, 0.0, 0.2, 0), + initial_plan_provider=provider, + ) + provider.assert_not_called() + + +def test_initial_plan_provider_cannot_bypass_phase_gate_validation() -> None: + engine, _ = _engine() + action = PhaseGateAction() + engine.register(action) + provider = Mock( + side_effect=lambda request, context: _candidate_plan(action, request, context) + ) + with pytest.raises(ValueError, match="missing segment 'commit'"): + engine.start( + (_phase_gate_invocation(engine),), + _context(0.0, 0.0, 0.2, 0), + initial_plan_provider=provider, + ) + assert action.plan_count == 0 + + +def test_initial_plan_provider_keeps_physical_effect_verification() -> None: + engine, _ = _engine() + action = EffectAction() + engine.register(action) + initial = _context(0.0, 0.0, 0.2, 0) + provider = Mock(side_effect=action._plan) + session = engine.start( + (_invocation(engine, skill_id=action.skill_id),), + initial, + initial_plan_provider=provider, + ) + session.tick(initial) + session.tick(_context(0.1, 0.0, 0.2, 0)) + waiting = session.tick(_context(0.2, 0.2, 0.2, 0)) + assert waiting.pending_effect is not None + assert waiting.status is ExecutionStatus.RUNNING + assert waiting.task_state.get_held_object("arm") is None + completed = session.tick( + _context(0.3, 0.2, 0.2, 0), + effect_result=_effect_result( + waiting.pending_effect.verification_id, + torch.tensor([True]), + torch.tensor([False]), + ), + ) + assert completed.status is ExecutionStatus.COMPLETED + assert completed.task_state.get_held_object("arm") is not None + provider.assert_called_once() + + @pytest.mark.parametrize( ("segment_name", "message"), (("missing", "missing segment"), ("prepare", "first trajectory segment")), diff --git a/tests/sim/atomic_actions/test_motion_strategy_e2e.py b/tests/sim/atomic_actions/test_motion_strategy_e2e.py index 38bff128c..a732cf208 100644 --- a/tests/sim/atomic_actions/test_motion_strategy_e2e.py +++ b/tests/sim/atomic_actions/test_motion_strategy_e2e.py @@ -23,7 +23,11 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.robots import CobotMagicCfg -from embodichain.lab.sim.planners import MotionGenerator, MotionGenCfg, ToppraPlannerCfg +from embodichain.lab.sim.motion.planners import ( + MotionGenerator, + MotionGenCfg, + ToppraPlannerCfg, +) from embodichain.lab.sim.atomic_actions import ( ActionInvocation, AtomicActionEngine, diff --git a/tests/sim/atomic_actions/test_trajectory_ops.py b/tests/sim/atomic_actions/test_trajectory_ops.py index 23b14a7af..b3e62e2be 100644 --- a/tests/sim/atomic_actions/test_trajectory_ops.py +++ b/tests/sim/atomic_actions/test_trajectory_ops.py @@ -39,7 +39,7 @@ split_three_segments, translate_pose_world, ) -from embodichain.lab.sim.planners import MoveType +from embodichain.lab.sim.motion.planners import MoveType CPU = torch.device("cpu") diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index 9366cff15..139b1be9b 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -34,7 +34,7 @@ ArticulationAffordanceGeometry, TimedTrajectory, ) -from embodichain.lab.sim.solvers import BaseSolver +from embodichain.lab.sim.motion.solvers import BaseSolver from scripts.tutorials.atomic_action.dynamic_obstacle_recovery import ( _animate_obstacle_to_pose, _blocking_obstacle_pose, @@ -543,7 +543,9 @@ def test_ur_tutorial_solver_limits_skip_noop_hard_limit_warning(factory) -> None upper_qpos_limits=solver_limits[:, 1].clone(), ) - with patch("embodichain.lab.sim.solvers.base_solver.logger.log_warning") as warning: + with patch( + "embodichain.lab.sim.motion.solvers.base_solver.logger.log_warning" + ) as warning: BaseSolver.update_with_robot_limit(solver, UR_RUNTIME_QPOS_LIMITS) assert not warning.called diff --git a/tests/sim/planners/test_base_planner.py b/tests/sim/motion/planners/test_base_planner.py similarity index 96% rename from tests/sim/planners/test_base_planner.py rename to tests/sim/motion/planners/test_base_planner.py index 0bf1d09c2..8ae69443f 100644 --- a/tests/sim/planners/test_base_planner.py +++ b/tests/sim/motion/planners/test_base_planner.py @@ -20,7 +20,7 @@ import pytest -from embodichain.lab.sim.planners.base_planner import CollisionWorldInfo +from embodichain.lab.sim.motion.planners.base_planner import CollisionWorldInfo def test_collision_world_info_represents_one_contract() -> None: diff --git a/tests/sim/planners/test_curobo_integration.py b/tests/sim/motion/planners/test_curobo_integration.py similarity index 98% rename from tests/sim/planners/test_curobo_integration.py rename to tests/sim/motion/planners/test_curobo_integration.py index 5c941900c..f543da248 100644 --- a/tests/sim/planners/test_curobo_integration.py +++ b/tests/sim/motion/planners/test_curobo_integration.py @@ -38,13 +38,13 @@ from embodichain.lab.sim.robots import FrankaPandaCfg # noqa: E402 from embodichain.lab.sim.cfg import RigidBodyAttributesCfg # noqa: E402 from embodichain.lab.sim.shapes import CubeCfg # noqa: E402 -from embodichain.lab.sim.planners import ( # noqa: E402 +from embodichain.lab.sim.motion.planners import ( # noqa: E402 MotionGenCfg, MotionGenOptions, MotionGenerator, PlanState, ) -from embodichain.lab.sim.planners.curobo.curobo_planner import ( # noqa: E402 +from embodichain.lab.sim.motion.planners.curobo.curobo_planner import ( # noqa: E402 CuroboPlanOptions, CuroboPlanner, CuroboPlannerCfg, diff --git a/tests/sim/planners/test_curobo_planner.py b/tests/sim/motion/planners/test_curobo_planner.py similarity index 86% rename from tests/sim/planners/test_curobo_planner.py rename to tests/sim/motion/planners/test_curobo_planner.py index 552ef6dfc..5d631d623 100644 --- a/tests/sim/planners/test_curobo_planner.py +++ b/tests/sim/motion/planners/test_curobo_planner.py @@ -33,8 +33,8 @@ import torch import yaml -from embodichain.lab.sim.planners import CuroboPlannerCfg -from embodichain.lab.sim.planners.curobo.curobo_planner import ( +from embodichain.lab.sim.motion.planners import CuroboPlannerCfg +from embodichain.lab.sim.motion.planners.curobo.curobo_planner import ( CuroboPlanOptions, CuroboPlanner, CuroboPlannerCfg as CuroboPlannerCfgDirect, @@ -46,7 +46,7 @@ _torch_cuda_graph_capture_mode, _validate_dynamic_obstacles, ) -from embodichain.lab.sim.planners.curobo.curobo_yaml import ( +from embodichain.lab.sim.motion.planners.curobo.curobo_yaml import ( _mesh_to_obstacle_entry, _parse_mimic_joint_names, generate_curobo_world_yaml, @@ -701,7 +701,10 @@ def test_dynamic_update_uses_registry_id_in_curobo_backend(): assert [(name, env_idx) for name, _, env_idx in updates] == [("registry_cube", 0)] -def test_validate_joint_trajectory_checks_every_exact_sample_in_curobo_order(): +@pytest.mark.parametrize("failing_cost", ["bound", "self", "scene"]) +def test_validate_joint_trajectory_checks_every_exact_sample_in_curobo_order( + failing_cost, +): """The collision gate preserves samples and maps simulator joint order.""" planner = object.__new__(CuroboPlanner) planner.cfg = CuroboPlannerCfg( @@ -716,12 +719,18 @@ def from_position(position, *, joint_names): joint_states.append((position.clone(), tuple(joint_names))) return SimpleNamespace(position=position) - def validate(sample, *, env_query_idx): - collision_queries.append((sample.clone(), env_query_idx.clone())) - validity = torch.ones(sample.shape[:2], dtype=torch.bool) - if len(collision_queries) == 2: - validity[1, 0] = False - return validity + def kinematics(sample): + return SimpleNamespace(position=sample, robot_spheres=sample.unsqueeze(-2)) + + def cost(kind, sample): + value = torch.zeros(sample.shape[:2]) + if kind == failing_cost and sample[1, 0, 0] == pytest.approx(3.1): + value[1, 0] = 1.0 + return value + + def scene_cost(state, *, idxs_env_query): + collision_queries.append((state.position.clone(), idxs_env_query.clone())) + return cost("scene", state.position) planner._bindings = SimpleNamespace( JointState=SimpleNamespace(from_position=from_position), @@ -729,7 +738,17 @@ def validate(sample, *, env_query_idx): backend = SimpleNamespace( sim_joint_names=["sim_left", "sim_right"], sim_to_curobo_col_idx=None, - collision_checker=SimpleNamespace(validate=validate), + collision_checker=SimpleNamespace( + get_kinematics=kinematics, + setup_batch_tensors=lambda batch, horizon: None, + self_collision_cost=object(), + get_bound=lambda sample: cost("bound", sample), + get_self_collision=lambda spheres: cost("self", spheres.squeeze(-2)), + collision_constraint=SimpleNamespace( + update_num_spheres=lambda count, batch, horizon: None, + forward=scene_cost, + ), + ), profile=SimpleNamespace( sim_to_curobo_joint_names={ "sim_left": "curobo_left", @@ -979,7 +998,7 @@ def _make_curobo_engine( from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, ) - from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator + from embodichain.lab.sim.motion.planners import MotionGenCfg, MotionGenerator motion_generator = MotionGenerator( MotionGenCfg( @@ -1094,3 +1113,105 @@ def test_curobo_uses_accelerator_with_cpu_physics(): finally: sim.destroy() SimulationManager.flush_cleanup_queue() + + +@pytest.mark.parametrize("use_current_fallback", [False, True]) +def test_robot_yaml_cache_invalidates_when_locked_joint_initial_state_changes( + tmp_path, monkeypatch, use_current_fallback +): + """A changed gripper pose must rebuild its static collision geometry.""" + import embodichain.lab.sim.motion.planners.curobo.curobo_yaml as yaml_module + + urdf = tmp_path / "robot.urdf" + urdf.write_text("") + planner = object.__new__(CuroboPlanner) + planner.cfg = CuroboPlannerCfg(robot_uid="cache-test") + planner.cfg.auto_gen.cache_dir = str(tmp_path) + planner._curobo_device = torch.device("cpu") + current = torch.tensor([[0.0, 0.02]]) + planner.robot = SimpleNamespace( + joint_names=("arm_joint", "finger_joint"), + control_parts={"arm": ["arm_joint"]}, + cfg=SimpleNamespace( + fpath=str(urdf), init_qpos=None if use_current_fallback else [0.0, 0.02] + ), + get_qpos=lambda: current.clone(), + ) + generated = [] + + def generate(robot, control_part, output_path, **kwargs): + from pathlib import Path + + generated.append(output_path) + Path(output_path).write_text("robot_cfg: {}") + return output_path + + monkeypatch.setattr(yaml_module, "generate_curobo_robot_yaml", generate) + original = planner._auto_generate_robot_yaml("arm", "tool") + assert planner._auto_generate_robot_yaml("arm", "tool") == original + assert len(generated) == 1 + if use_current_fallback: + current[0, 1] = 0.04 + else: + planner.robot.cfg.init_qpos[1] = 0.04 + changed = planner._auto_generate_robot_yaml("arm", "tool") + assert changed != original + assert len(generated) == 2 + if use_current_fallback: + current[0, 0] = 0.5 + else: + planner.robot.cfg.init_qpos[0] = 0.5 + assert planner._auto_generate_robot_yaml("arm", "tool") == changed + + +@pytest.mark.parametrize("use_current_fallback", [False, True]) +def test_runtime_backend_rejects_changed_locked_joints_until_explicit_close( + use_current_fallback, +): + """Runtime cache hits must obey the same lock signature as generated YAMLs.""" + planner = object.__new__(CuroboPlanner) + planner.cfg = CuroboPlannerCfg(robot_uid="runtime-lock", use_cuda_graph=False) + planner._curobo_device = torch.device("cpu") + current = torch.tensor([[0.0, 0.02]]) + planner.robot = SimpleNamespace( + joint_names=("arm_joint", "finger_joint"), + control_parts={"arm": ["arm_joint"]}, + cfg=SimpleNamespace(init_qpos=None if use_current_fallback else [0.0, 0.02]), + get_qpos=lambda: current.clone(), + ) + planner._backend_cache = {} + built, closed = [], [] + + def build(**kwargs): + backend = SimpleNamespace(planner=object(), use_cuda_graph=False) + built.append(backend) + return backend + + planner._materialize_profile = lambda control_part: object() + planner._resolve_sim_joint_names = lambda control_part: ["arm_joint"] + planner._build_backend = build + planner._warmup_backend = lambda backend: None + planner._close_planner = closed.append + + original = planner._get_backend("arm", 1) + assert planner._get_backend("arm", 1) is original + if use_current_fallback: + current[0, 0] = 0.5 + else: + planner.robot.cfg.init_qpos[0] = 0.5 + assert planner._get_backend("arm", 1) is original + if use_current_fallback: + current[0, 1] = 0.04 + else: + planner.robot.cfg.init_qpos[1] = 0.04 + with pytest.raises(RuntimeError, match="locked-joint configuration changed"): + planner._get_backend("arm", 1) + assert len(built) == 1 + assert not closed + + planner.close() + assert closed == [original.planner] + replacement = planner._get_backend("arm", 1) + assert replacement is not original + assert replacement.robot_lock_signature != original.robot_lock_signature + assert len(built) == 2 diff --git a/tests/sim/planners/test_motion_generator.py b/tests/sim/motion/planners/test_motion_generator.py similarity index 98% rename from tests/sim/planners/test_motion_generator.py rename to tests/sim/motion/planners/test_motion_generator.py index 7b6b7ed0d..1db063937 100644 --- a/tests/sim/planners/test_motion_generator.py +++ b/tests/sim/motion/planners/test_motion_generator.py @@ -24,8 +24,8 @@ from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.robots import CobotMagicCfg -from embodichain.lab.sim.planners.utils import TrajectorySampleMethod -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod +from embodichain.lab.sim.motion.planners import ( MotionGenerator, MotionGenCfg, MotionGenOptions, diff --git a/tests/sim/planners/test_motion_generator_batched.py b/tests/sim/motion/planners/test_motion_generator_batched.py similarity index 98% rename from tests/sim/planners/test_motion_generator_batched.py rename to tests/sim/motion/planners/test_motion_generator_batched.py index 4b84350d9..81137248d 100644 --- a/tests/sim/planners/test_motion_generator_batched.py +++ b/tests/sim/motion/planners/test_motion_generator_batched.py @@ -22,12 +22,15 @@ import pytest import torch -from embodichain.lab.sim.planners.base_planner import CollisionWorldInfo, PlanOptions -from embodichain.lab.sim.planners.motion_generator import ( +from embodichain.lab.sim.motion.planners.base_planner import ( + CollisionWorldInfo, + PlanOptions, +) +from embodichain.lab.sim.motion.planners.motion_generator import ( MotionGenerator, MotionGenOptions, ) -from embodichain.lab.sim.planners.utils import PlanState, PlanResult, MoveType +from embodichain.lab.sim.motion.planners.utils import PlanState, PlanResult, MoveType BATCH_SIZE = 2 CONTROLLED_DOF = 6 @@ -89,7 +92,7 @@ def resample( return torch.lerp(trajectory[:, lower], trajectory[:, upper], weights) monkeypatch.setattr( - "embodichain.lab.sim.planners.motion_generator.resample_with_distance", + "embodichain.lab.sim.motion.planners.motion_generator.resample_with_distance", resample, ) diff --git a/tests/sim/planners/test_neural_batched.py b/tests/sim/motion/planners/test_neural_batched.py similarity index 93% rename from tests/sim/planners/test_neural_batched.py rename to tests/sim/motion/planners/test_neural_batched.py index 78055af88..b7b73e791 100644 --- a/tests/sim/planners/test_neural_batched.py +++ b/tests/sim/motion/planners/test_neural_batched.py @@ -21,8 +21,8 @@ class TestNeuralParseWaypoints: def test_parse_waypoints_batched(self): - from embodichain.lab.sim.planners.neural_planner import NeuralPlanner - from embodichain.lab.sim.planners.utils import PlanState, MoveType + from embodichain.lab.sim.motion.planners.neural_planner import NeuralPlanner + from embodichain.lab.sim.motion.planners.utils import PlanState, MoveType # Build a minimal planner shell by stubbing __init__ planner = NeuralPlanner.__new__(NeuralPlanner) @@ -45,11 +45,11 @@ def test_parse_waypoints_batched(self): class TestNeuralPlanBatched: def test_plan_returns_batched_success(self, monkeypatch): - from embodichain.lab.sim.planners.neural_planner import ( + from embodichain.lab.sim.motion.planners.neural_planner import ( NeuralPlanner, NeuralPlanOptions, ) - from embodichain.lab.sim.planners.utils import PlanState, MoveType + from embodichain.lab.sim.motion.planners.utils import PlanState, MoveType planner = NeuralPlanner.__new__(NeuralPlanner) planner.device = torch.device("cpu") @@ -112,11 +112,11 @@ def compute_fk(self, qpos=None, name=None, to_matrix=True): class TestNeuralEarlyConvergenceHold: def test_converged_env_holds_qpos(self): - from embodichain.lab.sim.planners.neural_planner import ( + from embodichain.lab.sim.motion.planners.neural_planner import ( NeuralPlanner, NeuralPlanOptions, ) - from embodichain.lab.sim.planners.utils import PlanState, MoveType + from embodichain.lab.sim.motion.planners.utils import PlanState, MoveType planner = NeuralPlanner.__new__(NeuralPlanner) planner.device = torch.device("cpu") diff --git a/tests/sim/planners/test_neural_planner.py b/tests/sim/motion/planners/test_neural_planner.py similarity index 98% rename from tests/sim/planners/test_neural_planner.py rename to tests/sim/motion/planners/test_neural_planner.py index c1e2d2786..869052657 100644 --- a/tests/sim/planners/test_neural_planner.py +++ b/tests/sim/motion/planners/test_neural_planner.py @@ -18,7 +18,7 @@ import torch import pytest -from embodichain.lab.sim.planners import ( +from embodichain.lab.sim.motion.planners import ( MotionGenCfg, MotionGenOptions, MotionGenerator, @@ -27,7 +27,7 @@ NeuralPlannerCfg, PlanState, ) -from embodichain.lab.sim.planners.neural_planner import ( +from embodichain.lab.sim.motion.planners.neural_planner import ( NeuralPlanOptions, _WaypointTransformerActor, ) @@ -349,7 +349,7 @@ def test_safe_torch_load_roundtrip(tmp_path): path = tmp_path / "checkpoint.pt" torch.save(checkpoint, path) - from embodichain.lab.sim.planners.neural_planner import _safe_torch_load + from embodichain.lab.sim.motion.planners.neural_planner import _safe_torch_load loaded = _safe_torch_load(path, map_location=torch.device("cpu")) assert "agent" in loaded diff --git a/tests/sim/planners/test_plan_state_batched.py b/tests/sim/motion/planners/test_plan_state_batched.py similarity index 92% rename from tests/sim/planners/test_plan_state_batched.py rename to tests/sim/motion/planners/test_plan_state_batched.py index fc195655b..74c8e8b45 100644 --- a/tests/sim/planners/test_plan_state_batched.py +++ b/tests/sim/motion/planners/test_plan_state_batched.py @@ -19,7 +19,12 @@ import pytest import torch -from embodichain.lab.sim.planners.utils import PlanState, PlanResult, MoveType, MovePart +from embodichain.lab.sim.motion.planners.utils import ( + PlanState, + PlanResult, + MoveType, + MovePart, +) class TestPlanStateBatched: @@ -78,13 +83,13 @@ def test_rejects_inconsistent_B(self): PlanState.from_qpos(torch.zeros(2, 7)), PlanState.from_qpos(torch.zeros(3, 7)), ] - from embodichain.lab.sim.planners import base_planner as bp + from embodichain.lab.sim.motion.planners import base_planner as bp with pytest.raises(ValueError): bp._check_batch_consistency(states, expected_b=None, robot_num_instances=2) def test_rejects_expected_b_mismatch(self): - from embodichain.lab.sim.planners import base_planner as bp + from embodichain.lab.sim.motion.planners import base_planner as bp states = [ PlanState.from_qpos(torch.zeros(2, 7)), diff --git a/tests/sim/planners/test_toppra_batched.py b/tests/sim/motion/planners/test_toppra_batched.py similarity index 91% rename from tests/sim/planners/test_toppra_batched.py rename to tests/sim/motion/planners/test_toppra_batched.py index a769327ce..4cfcde346 100644 --- a/tests/sim/planners/test_toppra_batched.py +++ b/tests/sim/motion/planners/test_toppra_batched.py @@ -23,8 +23,8 @@ import pytest import torch -from embodichain.lab.sim.planners.toppra_planner import _toppra_solve_one_env -from embodichain.lab.sim.planners.utils import TrajectorySampleMethod +from embodichain.lab.sim.motion.planners.toppra_planner import _toppra_solve_one_env +from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod # Note: TOPPRA's ParametrizeConstAccel robustly avoids returning None for smooth @@ -105,7 +105,7 @@ def test_solve_one_env_duplicate_plateau(self): class TestToppraCfgFields: def test_cfg_defaults(self): - from embodichain.lab.sim.planners.toppra_planner import ToppraPlannerCfg + from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlannerCfg cfg = ToppraPlannerCfg(robot_uid="x") assert cfg.max_workers is None @@ -117,7 +117,7 @@ def test_mp_context_auto_resolution(self): # are honored as-is. import torch - from embodichain.lab.sim.planners.toppra_planner import ToppraPlanner + from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanner resolve = ToppraPlanner._resolve_mp_context # Auto: CPU -> fork, CUDA -> spawn (creating a torch.device does NOT @@ -132,7 +132,7 @@ def test_mp_context_auto_resolution(self): class TestToppraPlanBatched: def _make_planner(self): - from embodichain.lab.sim.planners.toppra_planner import ( + from embodichain.lab.sim.motion.planners.toppra_planner import ( ToppraPlanner, ToppraPlannerCfg, ) @@ -151,8 +151,11 @@ def _make_planner(self): return planner, sim def test_plan_batched_quantity_uniform_N(self): - from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod - from embodichain.lab.sim.planners.toppra_planner import ToppraPlanOptions + from embodichain.lab.sim.motion.planners.utils import ( + PlanState, + TrajectorySampleMethod, + ) + from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanOptions planner, sim = self._make_planner() try: @@ -181,8 +184,11 @@ def test_plan_batched_quantity_uniform_N(self): om.SimulationManager.flush_cleanup_queue() def test_plan_batched_time_tailpads(self): - from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod - from embodichain.lab.sim.planners.toppra_planner import ToppraPlanOptions + from embodichain.lab.sim.motion.planners.utils import ( + PlanState, + TrajectorySampleMethod, + ) + from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanOptions planner, sim = self._make_planner() try: @@ -234,8 +240,11 @@ def test_plan_batched_pool_path(self, mp_context): # Exercise the real ProcessPoolExecutor branch (max_workers=2, B=3) # under both start methods: 'fork' is the CPU default, 'spawn' is the # GPU fallback. Both must plan correctly and reap workers on GC. - from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod - from embodichain.lab.sim.planners.toppra_planner import ( + from embodichain.lab.sim.motion.planners.utils import ( + PlanState, + TrajectorySampleMethod, + ) + from embodichain.lab.sim.motion.planners.toppra_planner import ( ToppraPlanner, ToppraPlannerCfg, ToppraPlanOptions, @@ -293,8 +302,11 @@ def test_workers_reaped_on_gc(self, mp_context): # of an abandoned planner (e.g. between tests in a long-running # session), __del__ -> _shutdown_pool() must terminate workers # synchronously so none leak. - from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod - from embodichain.lab.sim.planners.toppra_planner import ( + from embodichain.lab.sim.motion.planners.utils import ( + PlanState, + TrajectorySampleMethod, + ) + from embodichain.lab.sim.motion.planners.toppra_planner import ( ToppraPlanner, ToppraPlannerCfg, ToppraPlanOptions, @@ -358,13 +370,16 @@ def test_workers_reaped_on_gc(self, mp_context): @pytest.mark.slow class TestToppraNumericalRegression: def test_batched_equals_inline_single(self): - from embodichain.lab.sim.planners.toppra_planner import ( + from embodichain.lab.sim.motion.planners.toppra_planner import ( _toppra_solve_one_env, ToppraPlanner, ToppraPlannerCfg, ToppraPlanOptions, ) - from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod + from embodichain.lab.sim.motion.planners.utils import ( + PlanState, + TrajectorySampleMethod, + ) from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.robots import CobotMagicCfg diff --git a/tests/sim/planners/test_toppra_planner.py b/tests/sim/motion/planners/test_toppra_planner.py similarity index 96% rename from tests/sim/planners/test_toppra_planner.py rename to tests/sim/motion/planners/test_toppra_planner.py index ce343cdb4..ad5823132 100644 --- a/tests/sim/planners/test_toppra_planner.py +++ b/tests/sim/motion/planners/test_toppra_planner.py @@ -20,12 +20,12 @@ import pytest import numpy as np -from embodichain.lab.sim.planners.toppra_planner import ( +from embodichain.lab.sim.motion.planners.toppra_planner import ( ToppraPlanner, ToppraPlannerCfg, ToppraPlanOptions, ) -from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod +from embodichain.lab.sim.motion.planners.utils import PlanState, TrajectorySampleMethod from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.robots import CobotMagicCfg from embodichain.lab.sim.cfg import RenderCfg diff --git a/tests/sim/solvers/test_base_solver.py b/tests/sim/motion/solvers/test_base_solver.py similarity index 96% rename from tests/sim/solvers/test_base_solver.py rename to tests/sim/motion/solvers/test_base_solver.py index f2192b27f..c92878a82 100644 --- a/tests/sim/solvers/test_base_solver.py +++ b/tests/sim/motion/solvers/test_base_solver.py @@ -19,7 +19,11 @@ import pytest from embodichain.lab.sim.cfg import RobotCfg -from embodichain.lab.sim.solvers import DifferentialSolverCfg, OPWSolverCfg, URSolverCfg +from embodichain.lab.sim.motion.solvers import ( + DifferentialSolverCfg, + OPWSolverCfg, + URSolverCfg, +) UR5_DH_PARAMETERS = { "d1": 0.089159, @@ -139,7 +143,7 @@ def test_solver_cfg_from_dict_preserves_unannotated_config_attributes(): def test_solver_cfg_from_dict_ignores_unknown_fields(monkeypatch): warnings = [] monkeypatch.setattr( - "embodichain.lab.sim.solvers.base_solver.logger.log_warning", + "embodichain.lab.sim.motion.solvers.base_solver.logger.log_warning", warnings.append, ) diff --git a/tests/sim/solvers/test_differential_solver.py b/tests/sim/motion/solvers/test_differential_solver.py similarity index 100% rename from tests/sim/solvers/test_differential_solver.py rename to tests/sim/motion/solvers/test_differential_solver.py diff --git a/tests/sim/solvers/test_neural_ik_solver.py b/tests/sim/motion/solvers/test_neural_ik_solver.py similarity index 97% rename from tests/sim/solvers/test_neural_ik_solver.py rename to tests/sim/motion/solvers/test_neural_ik_solver.py index 324c329e9..70af97f44 100644 --- a/tests/sim/solvers/test_neural_ik_solver.py +++ b/tests/sim/motion/solvers/test_neural_ik_solver.py @@ -22,8 +22,8 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.robots.franka_panda import FrankaPandaCfg -from embodichain.lab.sim.solvers import NeuralIKSolverCfg -from embodichain.lab.sim.solvers.neural_ik_solver import _build_mlp +from embodichain.lab.sim.motion.solvers import NeuralIKSolverCfg +from embodichain.lab.sim.motion.solvers.neural_ik_solver import _build_mlp from embodichain.utils.utility import reset_all_seeds NUM_ARM_JOINTS = 7 diff --git a/tests/sim/solvers/test_opw_solver.py b/tests/sim/motion/solvers/test_opw_solver.py similarity index 100% rename from tests/sim/solvers/test_opw_solver.py rename to tests/sim/motion/solvers/test_opw_solver.py diff --git a/tests/sim/solvers/test_pink_solver.py b/tests/sim/motion/solvers/test_pink_solver.py similarity index 99% rename from tests/sim/solvers/test_pink_solver.py rename to tests/sim/motion/solvers/test_pink_solver.py index 8a8a27721..134272c29 100644 --- a/tests/sim/solvers/test_pink_solver.py +++ b/tests/sim/motion/solvers/test_pink_solver.py @@ -30,8 +30,10 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import RobotCfg, RenderCfg from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim.solvers.null_space_posture_task import NullSpacePostureTask -from embodichain.lab.sim.solvers.pink_solver import PinkSolverCfg +from embodichain.lab.sim.motion.solvers.null_space_posture_task import ( + NullSpacePostureTask, +) +from embodichain.lab.sim.motion.solvers.pink_solver import PinkSolverCfg PLANAR_URDF = """ diff --git a/tests/sim/solvers/test_pinocchio_solver.py b/tests/sim/motion/solvers/test_pinocchio_solver.py similarity index 100% rename from tests/sim/solvers/test_pinocchio_solver.py rename to tests/sim/motion/solvers/test_pinocchio_solver.py diff --git a/tests/sim/solvers/test_pytorch_solver.py b/tests/sim/motion/solvers/test_pytorch_solver.py similarity index 100% rename from tests/sim/solvers/test_pytorch_solver.py rename to tests/sim/motion/solvers/test_pytorch_solver.py diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/motion/solvers/test_srs_solver.py similarity index 99% rename from tests/sim/solvers/test_srs_solver.py rename to tests/sim/motion/solvers/test_srs_solver.py index f760ed616..60dc394fd 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/motion/solvers/test_srs_solver.py @@ -36,7 +36,7 @@ DexforceW1ArmSide, DexforceW1Version, ) -from embodichain.lab.sim.solvers.srs_solver import SRSSolver, SRSSolverCfg +from embodichain.lab.sim.motion.solvers.srs_solver import SRSSolver, SRSSolverCfg class BaseSolverTest: diff --git a/tests/sim/solvers/test_ur_solver.py b/tests/sim/motion/solvers/test_ur_solver.py similarity index 99% rename from tests/sim/solvers/test_ur_solver.py rename to tests/sim/motion/solvers/test_ur_solver.py index 6fdcc8874..2d138c8d0 100644 --- a/tests/sim/solvers/test_ur_solver.py +++ b/tests/sim/motion/solvers/test_ur_solver.py @@ -23,7 +23,7 @@ from embodichain.data import get_data_path from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim.solvers import URSolverCfg +from embodichain.lab.sim.motion.solvers import URSolverCfg from embodichain.lab.sim.cfg import ( RenderCfg, JointDrivePropertiesCfg, diff --git a/tests/sim/motion/test_motion_imports.py b/tests/sim/motion/test_motion_imports.py new file mode 100644 index 000000000..b604ceb27 --- /dev/null +++ b/tests/sim/motion/test_motion_imports.py @@ -0,0 +1,67 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import subprocess +import sys + +import pytest + + +@pytest.mark.parametrize( + "first_module", + ["embodichain.lab.sim.objects", "embodichain.lab.sim.motion.planners"], +) +def test_motion_import_order_and_config_resolution(first_module: str) -> None: + script = """ +import importlib +import sys +import dexsim + +importlib.import_module(sys.argv[1]) +from embodichain.lab.sim import motion +from embodichain.lab.sim.cfg import RobotCfg +from embodichain.lab.sim.motion.solvers import SolverCfg, URSolverCfg +from embodichain.lab.sim.motion.workspace import RobotWorkspaceCfg + +solver = SolverCfg.from_dict({"class_type": "URSolver", "robot_type": "ur5"}) +robot = RobotCfg.from_dict({ + "solver_cfg": {"arm": {"class_type": "URSolver", "robot_type": "ur5"}}, + "workspace_cfg": {"arm": {"cache_path": "workspace.h5"}}, +}) +assert isinstance(solver, URSolverCfg) +assert isinstance(robot.solver_cfg["arm"], URSolverCfg) +assert isinstance(robot.workspace_cfg["arm"], RobotWorkspaceCfg) +for name in motion.__all__: + module = getattr(motion, name) + assert module is importlib.import_module("embodichain.lab.sim.motion." + name) +assert "solvers" in dir(motion) +try: + motion.unknown_motion_module +except AttributeError: + pass +else: + raise AssertionError("unknown motion attribute must raise AttributeError") +assert dexsim.get_world_num() == 0 +""" + result = subprocess.run( + [sys.executable, "-c", script, first_module], + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/sim/motion/trajectory_augmentation/__init__.py b/tests/sim/motion/trajectory_augmentation/__init__.py new file mode 100644 index 000000000..355d915ff --- /dev/null +++ b/tests/sim/motion/trajectory_augmentation/__init__.py @@ -0,0 +1,17 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations diff --git a/tests/sim/motion/trajectory_augmentation/test_cfg.py b/tests/sim/motion/trajectory_augmentation/test_cfg.py new file mode 100644 index 000000000..0dd12eb6e --- /dev/null +++ b/tests/sim/motion/trajectory_augmentation/test_cfg.py @@ -0,0 +1,175 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + TrajectoryAugmentationCfg, + TrajectoryGenerationJobCfg, +) + + +def _capabilities() -> dict[str, set[str]]: + return { + "source_ids": {"handwritten_qpos"}, + "validator_ids": {"task_success"}, + "profile_ids": { + "fixed_scene_initial_state", + "fixed_scene_tolerances", + "verified_motion", + "robot_execution_limits", + }, + "sink_ids": {"lerobot"}, + "operators": {"joint_residual", "retime"}, + } + + +def test_job_decodes_nested_schema_and_round_trips_without_imports() -> None: + cfg = TrajectoryGenerationJobCfg.from_mapping( + { + "augmentation": { + "seed": 17, + "factors": { + "timing": {"enabled": True, "duration_scales": [0.8, 1.0, 1.2]} + }, + }, + "collection": {"target_committed_episodes": 4}, + } + ) + cfg.validate_capabilities(**_capabilities()) + assert cfg.planning.batch_mode == "env_rows" + assert cfg.execution.pool_mode == "per_env_case" + assert cfg.persistence.sink == "lerobot" + assert cfg.augmentation.factors.timing.duration_scales == (0.8, 1.0, 1.2) + assert ( + TrajectoryGenerationJobCfg.from_mapping(cfg.to_dict()).to_dict() + == cfg.to_dict() + ) + + +@pytest.mark.parametrize( + "payload", + [ + {"seed": 4}, + {"source": {"callable": "os:system"}}, + {"execution": {"unknown": 1}}, + {"augmentation": {"factors": {"timing": {"unknown": 1}}}}, + {"collection": {"max_proposals": True}}, + {"persistence": {"async_write": "false"}}, + {"augmentation": {"seed": -1}}, + {"augmentation": {"seed": 2**63}}, + { + "augmentation": { + "factors": {"timing": {"enabled": True, "duration_scales": [0.0, 1.0]}} + } + }, + { + "augmentation": { + "factors": { + "timing": {"enabled": True, "duration_scales": [float("nan")]} + } + } + }, + {"augmentation": {"factors": {"timing": {"duration_scales": [2.0]}}}}, + {"augmentation": {"factors": {"spatial": {"joint_offset_scale": 2.0}}}}, + {"execution": {"ready_low_watermark": 16, "ready_high_watermark": 16}}, + {"execution": {"ready_max_bytes": 0}}, + {"collection": {"max_wall_time_s": float("inf")}}, + {"collection": {"target_committed_episodes": 101}}, + {"persistence": {"pending_max_bytes": 0}}, + ], +) +def test_decoder_rejects_unknown_fields_types_and_invalid_budgets( + payload: dict[str, object], +) -> None: + with pytest.raises(ValueError): + TrajectoryGenerationJobCfg.from_mapping(payload) + + +@pytest.mark.parametrize( + "payload", + [ + {"planning": {"batch_mode": "candidate_buckets"}}, + {"execution": {"pool_mode": "env_rows"}}, + {"collection": {"max_attempts_per_candidate": 2}}, + {"execution": {"pool_mode": "grouped_replicas"}}, + {"execution": {"scheduler": "slot_refill"}}, + {"execution": {"overlap_planning_and_physics": True}}, + {"reset": {"outer_mode": "host_reset"}}, + {"reset": {"on_initial_state_mismatch": "ignore"}}, + {"persistence": {"async_write": True}}, + {"validation": {"require_path_collision": False}}, + {"validation": {"require_task_success": False}}, + {"augmentation": {"factors": {"ik": {"enabled": True}}}}, + {"augmentation": {"factors": {"recovery": {"enabled": True}}}}, + ], +) +def test_unsupported_modes_cannot_be_silently_enabled( + payload: dict[str, object], +) -> None: + with pytest.raises(ValueError): + TrajectoryGenerationJobCfg.from_mapping(payload) + + +@pytest.mark.parametrize( + "registry", ["source_ids", "validator_ids", "profile_ids", "sink_ids"] +) +def test_preflight_requires_explicit_registered_implementations(registry: str) -> None: + capabilities = _capabilities() + capabilities[registry] = set() + with pytest.raises(ValueError, match="unregistered"): + TrajectoryGenerationJobCfg().validate_capabilities(**capabilities) + + +def test_spatial_and_timing_factors_require_available_operators() -> None: + cfg = TrajectoryGenerationJobCfg.from_mapping( + { + "augmentation": { + "factors": {"spatial": {"enabled": True, "method": "via_points"}} + } + } + ) + with pytest.raises(ValueError, match="via_points"): + cfg.validate_capabilities(**_capabilities()) + cfg = TrajectoryGenerationJobCfg.from_mapping( + { + "augmentation": { + "factors": {"timing": {"enabled": True, "duration_scales": [1.0, 2.0]}} + } + } + ) + capabilities = _capabilities() + capabilities["operators"] = set() + with pytest.raises(ValueError, match="retime"): + cfg.validate_capabilities(**capabilities) + + +def test_configuration_instances_do_not_share_nested_values() -> None: + first = TrajectoryGenerationJobCfg() + second = TrajectoryGenerationJobCfg() + first.augmentation.coverage.target_per_cell = 2 + assert second.augmentation.coverage.target_per_cell == 1 + + +def test_semantic_revalidation_rejects_mutated_configuration() -> None: + cfg = TrajectoryGenerationJobCfg() + cfg.persistence.async_write = True + with pytest.raises(ValueError, match="not implemented"): + cfg.validate_semantics() + with pytest.raises(ValueError, match="typed configuration"): + TrajectoryAugmentationCfg(factors={}) diff --git a/tests/sim/motion/trajectory_augmentation/test_contracts.py b/tests/sim/motion/trajectory_augmentation/test_contracts.py new file mode 100644 index 000000000..9fe2503a2 --- /dev/null +++ b/tests/sim/motion/trajectory_augmentation/test_contracts.py @@ -0,0 +1,301 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import replace + +import pytest +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation import ( + CandidateIdentity, + CandidateTrajectoryBatch, + CommitReceipt, + ExpertEpisode, + MotionSnapshot, + SceneCase, + TrajectoryPhase, + TrajectoryTemplate, + ValidationCheck, + ValidationResult, +) + + +def _case() -> SceneCase: + return SceneCase("case", "initial", "signature", "task", "robot") + + +def _identity(candidate_id: str = "candidate") -> CandidateIdentity: + return CandidateIdentity( + "case", "initial", candidate_id, "geometry", "source", "revision", "template" + ) + + +def _template(**changes: object) -> TrajectoryTemplate: + fields = dict( + source_id="source", + source_revision="revision", + template_id="template", + joint_names=("arm", "tool"), + positions=torch.tensor([[0.0, 0.0], [0.5, 0.0], [1.0, 0.0]]), + dt=torch.tensor([0.0, 0.1, 0.2]), + ) + fields.update(changes) + return TrajectoryTemplate(**fields) + + +def _batch(**changes: object) -> CandidateTrajectoryBatch: + fields = dict( + positions=torch.tensor( + [[[0.0, 0.0], [1.0, 0.0], [1.0, 0.0]], [[2.0, 0.0], [3.0, 0.0], [4.0, 0.0]]] + ), + dt=torch.tensor([[0.0, 0.1, 0.0], [0.0, 0.1, 0.2]]), + valid_length=torch.tensor([2, 3]), + identities=(_identity("candidate_0"), _identity("candidate_1")), + joint_names=("arm", "tool"), + source_row_indices=torch.tensor([0, 0]), + factors=({"duration_scale": 1.0}, {"duration_scale": 2.0}), + phases=( + (TrajectoryPhase("transit", 0, 2),), + (TrajectoryPhase("transit", 0, 3),), + ), + ) + fields.update(changes) + return CandidateTrajectoryBatch(**fields) + + +def _episode(**changes: object) -> ExpertEpisode: + fields = dict( + identity=_identity(), + observations={ + "joint_positions": torch.tensor([[0.0], [0.1], [0.2]]), + "image": torch.zeros((3, 2, 2, 3), dtype=torch.uint8), + }, + actions=torch.tensor([[0.1], [0.2]]), + timestamps=torch.tensor([0.0, 0.1, 0.2]), + action_representation="joint_position_target", + validation=ValidationResult((ValidationCheck("task_success", "passed"),)), + episode_id="episode", + commit_id="commit", + ) + fields.update(changes) + return ExpertEpisode(**fields) + + +def test_snapshot_copies_tensor_and_revision_inputs() -> None: + positions = torch.tensor([0.0, 1.0], requires_grad=True) + pose = torch.eye(4) + revisions = {"object": 1} + snapshot = MotionSnapshot( + _case(), + ("arm", "tool"), + positions, + torch.zeros(2), + pose, + {"object": pose}, + revisions, + ) + with torch.no_grad(): + positions[:] = 9 + pose[0, 3] = 5 + revisions["object"] = 2 + assert snapshot.joint_positions.tolist() == [0.0, 1.0] + assert not snapshot.joint_positions.requires_grad + assert snapshot.root_pose[0, 3] == 0 + assert snapshot.entity_poses["object"][0, 3] == 0 + assert snapshot.dependency_revisions["object"] == 1 + with pytest.raises(TypeError): + snapshot.dependency_revisions["object"] = 3 + + +@pytest.mark.parametrize( + "pose", + [ + torch.zeros((4, 4)), + torch.diag(torch.tensor([-1.0, 1.0, 1.0, 1.0])), + torch.eye(3), + torch.full((4, 4), float("nan")), + ], +) +def test_snapshot_rejects_invalid_or_reflected_se3(pose: torch.Tensor) -> None: + with pytest.raises(ValueError): + MotionSnapshot(_case(), ("arm",), torch.zeros(1), torch.zeros(1), pose) + + +def test_snapshot_rejects_incomplete_joint_state() -> None: + with pytest.raises(ValueError, match="complete joint_names"): + MotionSnapshot( + _case(), ("arm", "tool"), torch.zeros(1), torch.zeros(1), torch.eye(4) + ) + + +@pytest.mark.parametrize( + "changes", + [ + {"dt": torch.tensor([0.1, 0.1, 0.1])}, + {"dt": torch.tensor([0.0, 0.0, 0.1])}, + {"positions": torch.full((3, 2), float("inf"))}, + {"positions": torch.ones((3, 2), dtype=torch.int64)}, + {"joint_names": ("arm", "arm")}, + {"controlled_joint_indices": (2,)}, + {"controlled_joint_indices": (True,)}, + {"controlled_joint_indices": (0, 0)}, + {"representation": "eef_delta"}, + {"phases": (TrajectoryPhase("transit", 0, 4),)}, + {"phases": (TrajectoryPhase("a", 0, 2), TrajectoryPhase("b", 1, 3))}, + ], +) +def test_template_rejects_invalid_representation_timing_and_indices( + changes: dict[str, object], +) -> None: + with pytest.raises(ValueError): + _template(**changes) + + +def test_template_owns_inputs_and_leaves_unannotated_path_replay_only() -> None: + positions = torch.zeros((3, 2)) + dt = torch.tensor([0.0, 0.1, 0.1]) + template = _template(positions=positions, dt=dt) + positions.fill_(1) + dt.fill_(1) + assert template.positions.count_nonzero() == 0 + assert template.dt.tolist() == pytest.approx([0.0, 0.1, 0.1]) + assert template.phases == () + assert template.allowed_operators == () + + +def test_batch_preserves_repeated_source_rows_and_metadata_on_extraction() -> None: + batch = _batch() + assert batch.valid_mask.tolist() == [[True, True, False], [True, True, True]] + row = batch.row(1) + assert row.source_row_indices.tolist() == [0] + assert row.identities == (batch.identities[1],) + assert row.phases[0][0].stop_index == 3 + assert row.factors[0]["duration_scale"] == 2.0 + row.positions.fill_(9) + assert batch.positions[1, 0, 0] == 2 + + +def test_batch_accepts_empty_candidate_set() -> None: + batch = CandidateTrajectoryBatch( + torch.empty((0, 0, 2)), + torch.empty((0, 0)), + torch.empty((0,), dtype=torch.int64), + (), + ("arm", "tool"), + ) + assert batch.valid_mask.shape == (0, 0) + with pytest.raises(IndexError): + batch.row(0) + + +@pytest.mark.parametrize( + "changes", + [ + {"valid_length": torch.tensor([0, 3])}, + {"valid_length": torch.tensor([2, 4])}, + {"valid_length": torch.tensor([2.0, 3.0])}, + {"dt": torch.tensor([[0.0, 0.1, 0.1], [0.0, 0.1, 0.1]])}, + {"positions": torch.zeros((2, 3, 3))}, + {"identities": (_identity(), _identity())}, + {"factors": ({"scale": float("nan")}, {})}, + {"source_row_indices": torch.tensor([0, -1])}, + {"phases": ((TrajectoryPhase("transit", 0, 3),), ())}, + ], +) +def test_batch_rejects_invalid_lengths_padding_and_alignment( + changes: dict[str, object], +) -> None: + with pytest.raises(ValueError): + _batch(**changes) + + +def test_batch_rejects_padding_that_moves_the_robot() -> None: + positions = _batch().positions + positions[0, 2, 0] = 5 + with pytest.raises(ValueError, match="padded positions"): + _batch(positions=positions) + + +@pytest.mark.parametrize( + "status,accepted", + [("not_run", False), ("passed", True), ("failed", False), ("unavailable", False)], +) +def test_validation_accepts_only_completed_passing_checks( + status: str, accepted: bool +) -> None: + result = ValidationResult( + (ValidationCheck("collision", status, metrics={"clearance": 0.02}),) + ) + assert result.accepted is accepted + + +def test_validation_does_not_accept_empty_or_duplicate_checks() -> None: + assert not ValidationResult(()).accepted + with pytest.raises(ValueError, match="unique"): + ValidationResult( + ( + ValidationCheck("collision", "passed"), + ValidationCheck("collision", "passed"), + ) + ) + with pytest.raises(ValueError, match="finite"): + ValidationCheck("collision", "passed", metrics={"clearance": float("nan")}) + + +def test_episode_preserves_terminal_observation_and_copies_evidence() -> None: + observations = {"joint_positions": torch.tensor([[0.0], [0.1], [0.2]])} + metadata = {"lineage": ["parent"], "control_dt": 0.1} + episode = _episode(observations=observations, metadata=metadata) + observations["joint_positions"].fill_(9) + metadata["lineage"].append("unexpected") + assert episode.observations["joint_positions"][-1, 0] == pytest.approx(0.2) + assert episode.actions.shape[0] + 1 == episode.timestamps.shape[0] + assert episode.metadata["lineage"] == ("parent",) + + +@pytest.mark.parametrize( + "changes", + [ + {"observations": {"joint_positions": torch.zeros((2, 1))}}, + {"observations": {}}, + {"timestamps": torch.tensor([0.0, 0.1, 0.1])}, + {"actions": torch.zeros((0, 1))}, + {"metadata": {"host": object()}}, + {"metadata": {"metric": float("nan")}}, + {"phases": (TrajectoryPhase("transit", 0, 4),)}, + ], +) +def test_episode_rejects_noncausal_shapes_or_live_metadata( + changes: dict[str, object], +) -> None: + with pytest.raises(ValueError): + _episode(**changes) + + +def test_receipts_use_stable_ids_and_distinct_write_attempts() -> None: + receipt = CommitReceipt( + "episode", "candidate", 0, "dataset/episode", "commit", "case" + ) + retry = replace(receipt, submission_id=1, confirmed=False, error="storage offline") + assert receipt.commit_id == retry.commit_id + assert receipt.attempt_id == retry.attempt_id + assert receipt.submission_id != retry.submission_id + with pytest.raises(ValueError, match="only unconfirmed"): + replace(receipt, error="inconsistent") + with pytest.raises(ValueError, match="submission_id"): + replace(receipt, submission_id=-1) diff --git a/tests/sim/motion/trajectory_augmentation/test_coverage.py b/tests/sim/motion/trajectory_augmentation/test_coverage.py new file mode 100644 index 000000000..deb39824d --- /dev/null +++ b/tests/sim/motion/trajectory_augmentation/test_coverage.py @@ -0,0 +1,144 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation.contracts import TrajectoryPhase +from embodichain.lab.sim.motion.trajectory_augmentation.coverage import ( + CoverageIndex, + TrajectoryDescriptor, + describe_trajectory, +) + + +def descriptor(count: int = 3, duration: float = 1.0, offset: float = 0.0): + return describe_trajectory( + torch.linspace(0, 1, count)[:, None] + offset, + torch.linspace(0, duration, count), + torch.tensor([[0.0, 2.0]]), + samples_per_phase=8, + ) + + +def test_extra_samples_and_timing_do_not_change_geometry() -> None: + original, slow = descriptor(), descriptor(9, 2.0) + torch.testing.assert_close(original.geometry, slow.geometry) + torch.testing.assert_close(original.timing * 2, slow.timing) + index = CoverageIndex() + assert index.reserve("first", "family", original) + assert index.reserve("slow", "family", slow) + assert index.geometry_count == 0 + index.confirm("first") + index.confirm("slow") + assert index.geometry_count == 1 + + +def test_actual_near_duplicates_use_one_family_and_pending_quota() -> None: + index = CoverageIndex(target_per_geometry=2) + assert index.reserve("first", "one", descriptor()) + assert not index.reserve("copy", "two", descriptor(10)) + assert index.reserve("slow", "two", descriptor(10, 2.0)) + assert not index.reserve("third", "three", descriptor(10, 3.0)) + index.release("first") + assert index.reserve("third", "three", descriptor(10, 3.0)) + index.confirm("slow") + index.confirm("third") + assert index.geometry_count == 1 + + +def test_family_alias_survives_tracking_noise_and_failed_write() -> None: + index = CoverageIndex(geometry_tolerance=0.01) + assert index.reserve("first", "one", descriptor()) + assert index.reserve("slow", "two", descriptor(duration=2.0)) + index.release("first") + assert index.reserve("noisy", "two", descriptor(duration=3.0, offset=0.1)) + index.confirm("slow") + index.confirm("noisy") + assert index.geometry_count == 1 + + +def test_confirmation_is_idempotent_and_cannot_be_released() -> None: + index = CoverageIndex() + assert index.reserve("commit", "family", descriptor()) + index.confirm("commit") + index.confirm("commit") + assert index.geometry_count == 1 + with pytest.raises(ValueError, match="Committed"): + index.release("commit") + + +def test_retimed_phase_does_not_change_unlabelled_gap_identity() -> None: + limits = torch.tensor([[0.0, 2.0]]) + original = describe_trajectory( + torch.tensor([[0.0], [0.5], [1.0], [1.0]]), + torch.arange(4.0), + limits, + phases=(TrajectoryPhase("move", 0, 3),), + ) + slower = describe_trajectory( + torch.tensor([[0.0], [0.25], [0.5], [0.75], [1.0], [1.0]]), + torch.arange(6.0), + limits, + phases=(TrajectoryPhase("move", 0, 5),), + ) + assert original.phase_ids == slower.phase_ids + torch.testing.assert_close(original.geometry, slower.geometry) + index = CoverageIndex(target_per_geometry=1) + assert index.reserve("one", "first", original) + assert not index.reserve("two", "second", slower) + + +def test_phase_labels_are_owned_through_reservation() -> None: + labels = [["move", "free"]] + value = TrajectoryDescriptor(labels, torch.zeros(1, 2, 1), torch.zeros(1, 2)) + labels[0][0] = "changed" + assert value.phase_ids == (("move", "free"),) + + +def test_stationary_wait_counts_time_without_inventing_geometry() -> None: + result = describe_trajectory( + torch.zeros(4, 1), + torch.arange(4.0), + torch.tensor([[-1.0, 1.0]]), + ) + assert torch.equal(result.geometry, torch.full_like(result.geometry, 0.5)) + assert float(result.timing[0, -1]) == 3.0 + + +def test_phase_connections_and_unlabelled_samples_are_preserved() -> None: + positions = torch.tensor([[0.0], [0.0], [1.0], [1.0], [2.0]]) + result = describe_trajectory( + positions, + torch.arange(5.0), + torch.tensor([[0.0, 2.0]]), + phases=(TrajectoryPhase("first", 0, 2), TrajectoryPhase("second", 2, 4)), + samples_per_phase=4, + ) + assert result.geometry.shape == (3, 4, 1) + assert float(result.geometry[1, 0, 0]) == 0.0 + assert float(result.geometry[1, -1, 0]) == 0.5 + assert float(result.geometry[2, -1, 0]) == 1.0 + + +@pytest.mark.parametrize( + "times", [torch.tensor([0.0, 0.0]), torch.tensor([0.0, float("nan")])] +) +def test_invalid_actual_timestamps_are_rejected(times: torch.Tensor) -> None: + with pytest.raises(ValueError): + describe_trajectory(torch.zeros(2, 1), times, torch.tensor([[0.0, 1.0]])) diff --git a/tests/sim/motion/trajectory_augmentation/test_imports.py b/tests/sim/motion/trajectory_augmentation/test_imports.py new file mode 100644 index 000000000..5f9505afc --- /dev/null +++ b/tests/sim/motion/trajectory_augmentation/test_imports.py @@ -0,0 +1,33 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import subprocess +import sys + + +def test_public_core_import_uses_motion_namespace_without_creating_a_world() -> None: + script = """ +import dexsim +from embodichain.lab.sim.motion.trajectory_augmentation import TrajectoryGenerationJobCfg +TrajectoryGenerationJobCfg.from_mapping({}) +assert dexsim.get_world_num() == 0 +""" + completed = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=30 + ) + assert completed.returncode == 0, completed.stdout + completed.stderr diff --git a/tests/sim/motion/trajectory_augmentation/test_operators.py b/tests/sim/motion/trajectory_augmentation/test_operators.py new file mode 100644 index 000000000..65d89028e --- /dev/null +++ b/tests/sim/motion/trajectory_augmentation/test_operators.py @@ -0,0 +1,233 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import replace +import math + +import pytest +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation.contracts import ( + TrajectoryPhase, + TrajectoryTemplate, +) +from embodichain.lab.sim.motion.trajectory_augmentation.operators import ( + joint_residual, + retime, + rotate_grasp_about_object_axis, + validate_motion_limits, +) + + +def test_grasp_rotation_uses_object_origin_and_preserves_inputs() -> None: + object_pose = torch.eye(4) + object_pose[:3, 3] = torch.tensor([2.0, 3.0, 4.0]) + grasp = object_pose.clone() + grasp[0, 3] += 0.2 + grasp[2, 3] += 0.1 + before = grasp.clone() + candidates = rotate_grasp_about_object_axis( + object_pose, + grasp, + axis=torch.tensor([0.0, 0.0, 2.0]), + angles=torch.tensor([0.0, math.pi / 2]), + ) + torch.testing.assert_close(candidates[0], grasp) + torch.testing.assert_close(candidates[1, :3, 3], torch.tensor([2.0, 3.2, 4.1])) + torch.testing.assert_close( + candidates[1, :3, 0], torch.tensor([0.0, 1.0, 0.0]), atol=1e-6, rtol=0 + ) + torch.testing.assert_close(grasp, before) + candidates[0, 0, 3] = 100 + torch.testing.assert_close(grasp, before) + + +def test_grasp_variants_follow_rotated_object_frame() -> None: + # The object's local Z points along scene X. Quarter-turn augmentation + # must orbit that axis, even after translating the entire scene. + object_pose = torch.tensor( + [ + [0.0, 0.0, 1.0, 3.0], + [0.0, 1.0, 0.0, -2.0], + [-1.0, 0.0, 0.0, 1.0], + [0.0, 0.0, 0.0, 1.0], + ] + ) + grasp = object_pose.clone() + grasp[2, 3] -= 0.2 + candidate = rotate_grasp_about_object_axis( + object_pose, + grasp, + axis=torch.tensor([0.0, 0.0, 1.0]), + angles=torch.tensor([math.pi / 2]), + )[0] + torch.testing.assert_close(candidate[:3, 3], torch.tensor([3.0, -1.8, 1.0])) + torch.testing.assert_close(candidate[:3, :3].T @ candidate[:3, :3], torch.eye(3)) + assert torch.linalg.det(candidate[:3, :3]) == pytest.approx(1.0) + + +@pytest.mark.parametrize( + "bad", ["reflection", "zero_axis", "axis_shape", "nan_angle", "empty_angles"] +) +def test_grasp_rotation_rejects_invalid_geometry(bad: str) -> None: + pose = torch.eye(4) + axis = torch.tensor([0.0, 0.0, 1.0]) + angles = torch.zeros(1) + if bad == "reflection": + pose[0, 0] = -1 + elif bad == "zero_axis": + axis.zero_() + elif bad == "axis_shape": + axis = torch.ones(2) + elif bad == "nan_angle": + angles[0] = torch.nan + elif bad == "empty_angles": + angles = torch.zeros(0) + with pytest.raises(ValueError): + rotate_grasp_about_object_axis(pose, torch.eye(4), axis=axis, angles=angles) + + +def template() -> TrajectoryTemplate: + return TrajectoryTemplate( + source_id="reference", + source_revision="one", + template_id="pick", + joint_names=("arm", "gripper"), + positions=torch.tensor( + [[0.0, 0.0], [0.25, 0.0], [0.5, 0.0], [0.5, 0.0], [0.5, 0.1], [0.5, 0.1]] + ), + dt=torch.tensor([0.0, 0.1, 0.1, 0.1, 0.1, 0.1]), + phases=( + TrajectoryPhase( + "transit", 0, 3, allowed_operators=("retime", "joint_residual") + ), + TrajectoryPhase("grasp", 3, 6, kind="contact"), + ), + controlled_joint_indices=(0,), + allowed_operators=("retime", "joint_residual"), + ) + + +def test_retime_changes_length_preserves_contact_duration_and_endpoints() -> None: + source = template() + result = retime(source, duration_scale=2.0, control_dt=0.1) + assert result.positions.shape == (8, 2) + assert result.phases[0].stop_index == 5 + assert result.phases[1].start_index == 5 + torch.testing.assert_close(result.positions[4], source.positions[2]) + torch.testing.assert_close(result.positions[5:], source.positions[3:]) + torch.testing.assert_close(result.dt[6:].sum(), source.dt[4:].sum()) + torch.testing.assert_close(source.positions, template().positions) + + +def test_spatial_residual_preserves_locked_samples_and_uncontrolled_joints() -> None: + source = template() + generator = torch.Generator().manual_seed(5) + result = joint_residual( + source, + joint_limits=torch.tensor([[-1.0, 1.0], [-1.0, 1.0]]), + normalized_scale=0.05, + generator=generator, + ) + assert result.positions[1, 0] != source.positions[1, 0] + assert torch.equal( + result.positions[[0, 2, 3, 4, 5]], source.positions[[0, 2, 3, 4, 5]] + ) + assert torch.equal(result.positions[:, 1], source.positions[:, 1]) + + +def test_explicit_generator_is_reproducible_without_consuming_global_rng() -> None: + before = torch.random.get_rng_state().clone() + kwargs = dict( + joint_limits=torch.tensor([[-1.0, 1.0], [-1.0, 1.0]]), normalized_scale=0.05 + ) + first = joint_residual( + template(), generator=torch.Generator().manual_seed(2), **kwargs + ) + second = joint_residual( + template(), generator=torch.Generator().manual_seed(2), **kwargs + ) + assert torch.equal(first.positions, second.positions) + assert torch.equal(before, torch.random.get_rng_state()) + + +def test_unlabelled_qpos_cannot_be_augmented() -> None: + with pytest.raises(ValueError, match="annotated"): + retime(replace(template(), phases=()), duration_scale=2, control_dt=0.1) + + +def test_retime_does_not_silently_move_uncontrolled_schedule() -> None: + source = template() + source.positions[1, 1] = 0.01 + with pytest.raises(ValueError, match="uncontrolled"): + retime(source, duration_scale=2, control_dt=0.1) + + +def test_retime_rejects_clock_mismatch_and_unbounded_allocation() -> None: + with pytest.raises(ValueError, match="clock"): + retime(template(), duration_scale=2, control_dt=0.07) + with pytest.raises(ValueError, match="budget"): + retime(template(), duration_scale=1000, control_dt=0.1, max_samples=10) + + +def test_dynamic_limits_reject_overfast_motion_but_do_not_claim_task_success() -> None: + source = template() + failed = validate_motion_limits( + source, + velocity_limits=torch.tensor([1.0, 2.0]), + acceleration_limits=torch.full((2,), 100.0), + ) + assert not failed.accepted + slow = retime(source, duration_scale=4, control_dt=0.1) + passed = validate_motion_limits( + slow, + velocity_limits=torch.tensor([1.0, 2.0]), + acceleration_limits=torch.full((2,), 100.0), + ) + assert passed.accepted + assert passed.checks[0].check_id == "motion_limits" + + +def test_contact_internal_event_cannot_be_removed_by_host_clock() -> None: + source = replace(template(), dt=torch.tensor([0.0, 0.1, 0.1, 0.1, 0.05, 0.15])) + with pytest.raises(ValueError, match="clock"): + retime(source, duration_scale=2, control_dt=0.1) + + +def test_large_derivatives_return_a_failed_check() -> None: + source = replace( + template(), positions=template().positions * 1e38, dt=template().dt * 1e-3 + ) + result = validate_motion_limits( + source, velocity_limits=torch.ones(2), acceleration_limits=torch.ones(2) + ) + assert not result.accepted + assert result.checks[0].metrics["speed_ratio"] > 1e30 + + +def test_half_precision_derivatives_are_accumulated_in_high_precision() -> None: + source = replace( + template(), positions=template().positions.half(), dt=template().dt.half() + ) + result = validate_motion_limits( + source, + velocity_limits=torch.full((2,), 20.0), + acceleration_limits=torch.full((2,), 0.0001), + ) + assert not result.accepted + assert result.checks[0].metrics["acceleration_ratio"] > 65504 diff --git a/tests/sim/motion/trajectory_augmentation/test_session.py b/tests/sim/motion/trajectory_augmentation/test_session.py new file mode 100644 index 000000000..5162d4940 --- /dev/null +++ b/tests/sim/motion/trajectory_augmentation/test_session.py @@ -0,0 +1,572 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""CPU lifecycle checks for fixed-scene collection accounting.""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest +import torch + +from embodichain.lab.sim.motion.trajectory_augmentation.cfg import ( + TrajectoryGenerationJobCfg, +) +from embodichain.lab.sim.motion.trajectory_augmentation.contracts import ( + CandidateIdentity, + CandidateTrajectoryBatch, + CommitReceipt, + ExpertEpisode, + SceneCase, + TrajectoryPhase, + ValidationCheck, + ValidationResult, +) +from embodichain.lab.sim.motion.trajectory_augmentation.session import GenerationSession + +CASE = SceneCase("case", "initial", "scene_v1", "lift", "robot") +JOINT_NAMES = ("joint_a", "joint_b") +LIMITS = torch.tensor([[-1.0, 1.0], [-1.0, 1.0]]) +PLAN_VALID = ValidationResult((ValidationCheck("path_collision", "passed"),)) +ROLLOUT_VALID = ValidationResult( + ( + ValidationCheck("path_collision", "passed"), + ValidationCheck("task_success", "passed"), + ) +) + + +def _session(**overrides: object) -> GenerationSession: + cfg = TrajectoryGenerationJobCfg.from_mapping(overrides) + session = GenerationSession(cfg) + session.register_case(CASE, LIMITS, joint_names=JOINT_NAMES) + return session + + +def _propose(session: GenerationSession, case: SceneCase = CASE, **kwargs: object): + return session.propose( + case.scene_case_id, + case.initial_state_id, + source_id="handwritten_qpos", + source_revision="v1", + template_id="reference_0", + operator_id="joint_residual", + **kwargs, + ) + + +def _batch(*identities: CandidateIdentity) -> CandidateTrajectoryBatch: + # Two controlled joints and two waypoints; the host owns collision evidence. + count = len(identities) + return CandidateTrajectoryBatch( + torch.zeros(count, 2, 2), + torch.tensor([[0.0, 1.0]]).repeat(count, 1), + torch.full((count,), 2, dtype=torch.int64), + identities, + ("joint_a", "joint_b"), + source_row_indices=torch.arange(count), + ) + + +def _episode( + session: GenerationSession, + identity: CandidateIdentity, + *, + position: float = 0.5, + duration: float = 1.0, + validation: ValidationResult = ROLLOUT_VALID, +) -> ExpertEpisode: + episode_id, commit_id = session.episode_ids(identity) + return ExpertEpisode( + identity, + {"joint_positions": torch.tensor([[0.0, 0.0], [position, 0.0]])}, + torch.tensor([[position, 0.0]]), + torch.tensor([0.0, duration]), + "qpos", + validation, + episode_id, + commit_id, + phases=(TrajectoryPhase("move", 0, 2),), + ) + + +def _receipt(episode: ExpertEpisode, **kwargs: object) -> CommitReceipt: + return CommitReceipt( + episode.episode_id, + episode.identity.candidate_id, + episode.identity.attempt_id, + "memory:episode", + episode.commit_id, + episode.identity.scene_case_id, + **kwargs, + ) + + +def _start( + session: GenerationSession, + case: SceneCase = CASE, + *, + episode_byte_budget: int | None = 4096, + **kwargs: object, +): + identity, _ = _propose(session, case, **kwargs) + session.add_planned(_batch(identity), PLAN_VALID) + row = session.take_ready( + case.scene_case_id, + case.initial_state_id, + episode_byte_budget=episode_byte_budget, + ) + assert row.identities == (identity,) + session.mark_rollout_started(identity) + return identity + + +def test_proposals_use_case_local_ordinals_and_do_not_consume_global_rng() -> None: + first, second = _session(), _session() + other = replace(CASE, scene_case_id="other") + second.register_case(other, LIMITS, joint_names=JOINT_NAMES) + rng_before = torch.random.get_rng_state().clone() + expected = [_propose(first) for _ in range(3)] + actual = [] + for _ in range(3): + _propose(second, other) + actual.append(_propose(second)) + assert torch.equal(torch.random.get_rng_state(), rng_before) + for (left, left_rng), (right, right_rng) in zip(expected, actual): + assert left == right + assert torch.equal( + torch.rand(5, generator=left_rng), torch.rand(5, generator=right_rng) + ) + + +@pytest.mark.parametrize( + "limit", [{"ready_high_watermark": 1}, {"ready_max_bytes": 47}] +) +def test_ready_pool_rejects_whole_batch_when_count_or_bytes_would_overflow( + limit, +) -> None: + session = _session(execution={"ready_low_watermark": 0, **limit}) + identities = tuple(_propose(session)[0] for _ in range(2)) + with pytest.raises(BufferError, match="pool is full"): + session.add_planned(_batch(*identities), PLAN_VALID) + counts = session.snapshot()["counts"] + assert counts["planned_valid"] == counts["ready"] == 0 + assert counts["proposed"] == 2 + + +def test_planning_failure_and_assigned_cancellation_do_not_count_rollouts() -> None: + session = _session() + failed, _ = _propose(session) + with pytest.raises(ValueError, match="mandatory planning checks"): + session.add_planned(_batch(failed), ValidationResult(())) + session.release(failed) + ready, _ = _propose(session) + session.add_planned(_batch(ready), PLAN_VALID) + assert session.take_ready("case", "initial") is not None + session.release(ready) + assert session.snapshot()["counts"]["rollout_attempted"] == 0 + + +def test_failed_validation_statuses_and_release_reason_survive_payload_discard() -> ( + None +): + session = _session() + planning, _ = _propose(session) + unavailable = ValidationResult( + (ValidationCheck("path_collision", "unavailable", "large backend trace"),) + ) + with pytest.raises(ValueError, match="mandatory planning checks"): + session.add_planned(_batch(planning), unavailable) + session.release(planning, reason="collision provider missing") + rollout = _start(session) + failed = ValidationResult( + ( + ValidationCheck("path_collision", "passed"), + ValidationCheck("task_success", "failed"), + ) + ) + assert not session.accept_episode(_episode(session, rollout, validation=failed)) + diagnostics = session.snapshot()["diagnostics"] + assert diagnostics[planning.candidate_id]["reason"] == "collision provider missing" + assert diagnostics[planning.candidate_id]["planning_validation"] == unavailable + assert diagnostics[rollout.candidate_id]["reason"] == "rollout_validation_failed" + assert diagnostics[rollout.candidate_id]["rollout_validation"] == failed + with pytest.raises(TypeError): + diagnostics[planning.candidate_id]["reason"] = "changed" + + +def test_tail_batch_reserves_global_target_before_first_commands() -> None: + session = _session(collection={"target_committed_episodes": 2}) + identities = tuple(_propose(session)[0] for _ in range(3)) + session.add_planned(_batch(*identities), PLAN_VALID) + first = session.take_ready("case", "initial", episode_byte_budget=4096) + second = session.take_ready("case", "initial", episode_byte_budget=4096) + assert first.identities == (identities[0],) + assert second.identities == (identities[1],) + assert session.take_ready("case", "initial") is None + assert session.snapshot()["counts"]["rollout_attempted"] == 0 + session.release(identities[0]) + assert session.take_ready( + "case", "initial", episode_byte_budget=4096 + ).identities == (identities[2],) + + +def test_assignment_also_reserves_remaining_rollout_attempts() -> None: + session = _session( + collection={"target_committed_episodes": 2, "max_rollout_attempts": 2} + ) + failed = _start(session) + session.release(failed) + identities = tuple(_propose(session)[0] for _ in range(2)) + session.add_planned(_batch(*identities), PLAN_VALID) + assert session.take_ready("case", "initial") is not None + assert session.take_ready("case", "initial") is None + session.mark_rollout_started(identities[0]) + session.release(identities[0]) + assert session.stop_reason == "rollout_budget_exhausted" + + +def test_candidate_rows_are_owned_and_keep_source_mapping() -> None: + session = _session() + identity, _ = _propose(session) + batch = _batch(identity) + session.add_planned(batch, PLAN_VALID) + batch.positions.fill_(99) + ready = session.take_ready("case", "initial") + assert torch.equal(ready.positions, torch.zeros(1, 2, 2)) + assert ready.source_row_indices.tolist() == [0] + + +def test_candidate_joint_order_must_match_registered_limit_order() -> None: + session = _session() + identity, _ = _propose(session) + swapped = replace(_batch(identity), joint_names=tuple(reversed(JOINT_NAMES))) + with pytest.raises(ValueError, match="joint layout"): + session.add_planned(swapped, PLAN_VALID) + assert session.snapshot()["counts"]["planned_valid"] == 0 + + +@pytest.mark.parametrize( + "changed", + [{"action_representation": "qvel"}, {"actions": torch.zeros(1, 1)}], +) +def test_unsupported_action_labels_cannot_enter_expert_dataset(changed) -> None: + session = _session() + episode = _episode(session, _start(session)) + with pytest.raises(ValueError, match="full-joint qpos"): + session.accept_episode(replace(episode, **changed)) + assert session.snapshot()["counts"]["validated_accepted"] == 0 + + +def test_large_factor_strings_cannot_bypass_ready_byte_limit() -> None: + session = _session(execution={"ready_max_bytes": 4096}) + identity, _ = _propose(session) + oversized = replace(_batch(identity), factors=({"description": "x" * 5000},)) + with pytest.raises(BufferError, match="pool is full"): + session.add_planned(oversized, PLAN_VALID) + assert session.snapshot()["counts"]["ready"] == 0 + + +def test_large_episode_metadata_cannot_bypass_pending_byte_limit() -> None: + session = _session(persistence={"pending_max_bytes": 4096}) + episode = _episode(session, _start(session)) + oversized = replace(episode, metadata={"nested": {"description": "x" * 5000}}) + with pytest.raises(BufferError, match="reserved episode byte budget"): + session.accept_episode(oversized) + assert session.snapshot()["counts"]["pending_write"] == 0 + + +def test_only_confirmed_receipts_add_committed_count_and_formal_coverage() -> None: + session = _session() + identity = _start(session) + episode = _episode(session, identity) + assert session.accept_episode(episode) + before = session.snapshot() + assert before["counts"]["rollout_attempted"] == 1 + assert before["counts"]["validated_accepted"] == 1 + assert before["counts"]["pending_write"] == 1 + assert before["counts"]["committed"] == before["coverage"]["case"] == 0 + receipt = _receipt(episode) + assert session.apply_receipt(receipt) + assert not session.apply_receipt(receipt) + after = session.snapshot() + assert after["counts"]["committed"] == after["coverage"]["case"] == 1 + assert after["counts"]["pending_write"] == 0 + assert session.stop_reason == "target_reached" + assert before["counts"]["committed"] == 0 + with pytest.raises(TypeError): + after["counts"]["committed"] = 9 + + +@pytest.mark.parametrize( + "checks", + [ + (), + (ValidationCheck("task_success", "passed"),), + (ValidationCheck("path_collision", "passed"),), + ( + ValidationCheck("path_collision", "unavailable"), + ValidationCheck("task_success", "passed"), + ), + ], +) +def test_missing_or_unavailable_required_validation_cannot_accept(checks) -> None: + session = _session() + identity = _start(session) + assert not session.accept_episode( + _episode(session, identity, validation=ValidationResult(checks)) + ) + assert session.snapshot()["counts"]["validated_accepted"] == 0 + assert session.snapshot()["counts"]["running"] == 0 + + +def test_episode_and_receipt_must_match_original_candidate_identity() -> None: + session = _session() + identity = _start(session) + episode = _episode(session, identity) + with pytest.raises(ValueError, match="Episode and commit IDs"): + session.accept_episode(replace(episode, episode_id="foreign")) + assert session.accept_episode(episode) + with pytest.raises(ValueError, match="original candidate"): + session.apply_receipt(replace(_receipt(episode), scene_case_id="other")) + assert session.snapshot()["counts"]["pending_write"] == 1 + + +def test_failed_submission_retries_preserve_counts_and_ignore_old_duplicate_receipts() -> ( + None +): + session = _session() + identity = _start(session) + episode = _episode(session, identity) + session.accept_episode(episode) + failure = _receipt(episode, confirmed=False, error="disk unavailable") + session.apply_receipt(failure) + assert session.snapshot()["counts"]["pending_write"] == 0 + payload, submission_id = session.retry_write(episode.commit_id) + assert payload.commit_id == episode.commit_id + assert submission_id == 1 + assert not session.apply_receipt(failure) + assert session.snapshot()["counts"]["pending_write"] == 1 + with pytest.raises(ValueError, match="Contradictory final"): + session.apply_receipt(_receipt(episode)) + assert session.apply_receipt(_receipt(episode, submission_id=submission_id)) + counts = session.snapshot()["counts"] + assert ( + counts["rollout_attempted"] + == counts["validated_accepted"] + == counts["committed"] + == 1 + ) + + +def test_failed_write_cannot_retake_target_reserved_by_another_candidate() -> None: + session = _session() + episode = _episode(session, _start(session)) + session.accept_episode(episode) + session.apply_receipt(_receipt(episode, confirmed=False)) + _start(session) + with pytest.raises(RuntimeError, match="No collection quota"): + session.retry_write(episode.commit_id) + + +def test_actual_geometry_controls_duplicates_and_timing_variants_share_coverage() -> ( + None +): + session = _session( + collection={"target_committed_episodes": 3}, + augmentation={"coverage": {"target_per_cell": 2}}, + ) + first = _start(session) + first_episode = _episode(session, first) + assert session.accept_episode(first_episode) + # Proposed trajectories are identical zeros; real evidence decides coverage. + distinct = _start(session) + distinct_episode = _episode(session, distinct, position=-0.5) + assert session.accept_episode(distinct_episode) + session.apply_receipt(_receipt(distinct_episode)) + duplicate = _start(session) + assert not session.accept_episode(_episode(session, duplicate)) + timing = _start( + session, + geometry_family_id=first.geometry_family_id, + parent_id=first.candidate_id, + ) + timing_episode = _episode(session, timing, duration=2.0) + assert session.accept_episode(timing_episode) + session.apply_receipt(_receipt(first_episode)) + session.apply_receipt(_receipt(timing_episode)) + assert session.snapshot()["counts"]["committed"] == 3 + assert session.snapshot()["coverage"]["case"] == 2 + + +def test_same_measured_trajectory_has_independent_coverage_in_different_cases() -> None: + session = _session(collection={"target_committed_episodes": 2}) + other = replace(CASE, scene_case_id="other") + session.register_case(other, LIMITS, joint_names=JOINT_NAMES) + first = _episode(session, _start(session)) + second = _episode(session, _start(session, other)) + assert session.accept_episode(first) + assert session.accept_episode(second) + # Confirmation arrives after another case was assigned, in reverse order. + session.apply_receipt(_receipt(second)) + session.apply_receipt(_receipt(first)) + assert session.snapshot()["coverage"] == {"case": 1, "other": 1} + + +def test_pending_count_backpressure_prevents_rollout_and_retains_failed_writes() -> ( + None +): + session = _session( + collection={"target_committed_episodes": 2}, + persistence={"pending_episode_limit": 1}, + ) + identities = tuple(_propose(session)[0] for _ in range(2)) + session.add_planned(_batch(*identities), PLAN_VALID) + session.take_ready("case", "initial", episode_byte_budget=4096) + assert session.take_ready("case", "initial", episode_byte_budget=4096) is None + session.mark_rollout_started(identities[0]) + assert session.take_ready("case", "initial", episode_byte_budget=4096) is None + first = _episode(session, identities[0]) + assert session.accept_episode(first) + session.apply_receipt(_receipt(first, confirmed=False)) + assert session.take_ready("case", "initial", episode_byte_budget=4096) is None + assert session.snapshot()["counts"]["rollout_attempted"] == 1 + session.release(identities[0], reason="write_aborted") + assert session.take_ready( + "case", "initial", episode_byte_budget=4096 + ).identities == (identities[1],) + + +def test_assigned_byte_reservations_block_extra_rows_and_release_without_commands() -> ( + None +): + session = _session( + collection={"target_committed_episodes": 3}, + persistence={"pending_max_bytes": 8192}, + ) + identities = tuple(_propose(session)[0] for _ in range(3)) + session.add_planned(_batch(*identities), PLAN_VALID) + session.take_ready("case", "initial", episode_byte_budget=4096) + session.take_ready("case", "initial", episode_byte_budget=4096) + assert session.snapshot()["pending_reserved_bytes"] == 8192 + assert session.take_ready("case", "initial", episode_byte_budget=1) is None + session.release(identities[0]) + assert session.snapshot()["pending_reserved_bytes"] == 4096 + assert session.take_ready( + "case", "initial", episode_byte_budget=4096 + ).identities == (identities[2],) + assert session.snapshot()["counts"]["rollout_attempted"] == 0 + + +def test_unspecified_episode_byte_budget_conservatively_serializes_until_confirmation() -> ( + None +): + session = _session( + collection={"target_committed_episodes": 2}, + persistence={"pending_max_bytes": 8192}, + ) + identities = tuple(_propose(session)[0] for _ in range(2)) + session.add_planned(_batch(*identities), PLAN_VALID) + session.take_ready("case", "initial") + assert session.snapshot()["pending_reserved_bytes"] == 8192 + assert session.take_ready("case", "initial") is None + session.mark_rollout_started(identities[0]) + first = _episode(session, identities[0]) + assert session.accept_episode(first) + assert 0 < session.snapshot()["pending_reserved_bytes"] < 8192 + assert session.take_ready("case", "initial") is None + session.apply_receipt(_receipt(first)) + assert session.snapshot()["pending_reserved_bytes"] == 0 + assert session.take_ready("case", "initial").identities == (identities[1],) + + +def test_episode_exceeding_declared_byte_budget_releases_reserved_capacity() -> None: + session = _session() + identity = _start(session, episode_byte_budget=31) + assert session.snapshot()["pending_reserved_bytes"] == 31 + with pytest.raises(BufferError, match="reserved episode byte budget"): + session.accept_episode(_episode(session, identity)) + assert session.snapshot()["pending_reserved_bytes"] == 0 + assert session.snapshot()["counts"]["pending_write"] == 0 + assert ( + session.snapshot()["diagnostics"][identity.candidate_id]["reason"] + == "episode_byte_budget_exceeded" + ) + + +@pytest.mark.parametrize("budget", [0, -1, True, 1.5, 8193]) +def test_invalid_episode_byte_budget_does_not_assign_a_row(budget) -> None: + session = _session(persistence={"pending_max_bytes": 8192}) + identity, _ = _propose(session) + session.add_planned(_batch(identity), PLAN_VALID) + with pytest.raises(ValueError, match="episode_byte_budget"): + session.take_ready("case", "initial", episode_byte_budget=budget) + assert session.snapshot()["counts"]["ready"] == 1 + assert session.snapshot()["pending_reserved_bytes"] == 0 + + +def test_diagnostics_keep_bounded_failure_metrics_and_detail() -> None: + session = _session() + identity, _ = _propose(session) + checks = ValidationResult( + ( + ValidationCheck( + "path_collision", + "failed", + "x" * 1024, + {f"metric_{index}": float(index) for index in range(20)}, + ), + ) + ) + with pytest.raises(ValueError, match="mandatory planning checks"): + session.add_planned(_batch(identity), checks) + summary = session.snapshot()["diagnostics"][identity.candidate_id][ + "planning_validation" + ].checks[0] + assert summary.status == "failed" + assert summary.detail == "x" * 256 + assert len(summary.metrics) == 16 + assert summary.metrics["metric_15"] == 15.0 + + +def test_wall_time_uses_injected_clock_and_still_allows_final_receipt() -> None: + now = [0.0] + session = GenerationSession(TrajectoryGenerationJobCfg(), clock=lambda: now[0]) + session.register_case(CASE, LIMITS, joint_names=JOINT_NAMES) + episode = _episode(session, _start(session)) + assert session.accept_episode(episode) + now[0] = 61.0 + assert session.stop_reason == "wall_time_exhausted" + assert session.take_ready("case", "initial") is None + with pytest.raises(RuntimeError, match="budget"): + _propose(session) + session.apply_receipt(_receipt(episode)) + assert session.snapshot()["counts"]["committed"] == 1 + + +def test_proposal_budget_exhaustion_and_case_conditions_do_not_reset_history() -> None: + session = _session(collection={"max_proposals": 1}) + identity, _ = _propose(session) + session.release(identity) + session.register_case(CASE, LIMITS, joint_names=JOINT_NAMES) + assert session.stop_reason == "proposal_budget_exhausted" + with pytest.raises(RuntimeError, match="budget"): + _propose(session) + with pytest.raises(ValueError, match="fixed conditions"): + session.register_case( + replace(CASE, scene_signature="changed"), LIMITS, joint_names=JOINT_NAMES + ) diff --git a/tests/sim/workspace/__init__.py b/tests/sim/motion/workspace/__init__.py similarity index 100% rename from tests/sim/workspace/__init__.py rename to tests/sim/motion/workspace/__init__.py diff --git a/tests/sim/workspace/test_analyzer.py b/tests/sim/motion/workspace/test_analyzer.py similarity index 97% rename from tests/sim/workspace/test_analyzer.py rename to tests/sim/motion/workspace/test_analyzer.py index fd9f227fd..f5dace3fc 100644 --- a/tests/sim/workspace/test_analyzer.py +++ b/tests/sim/motion/workspace/test_analyzer.py @@ -23,7 +23,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.robots import CobotMagicCfg -from embodichain.lab.sim.workspace.configs import ( +from embodichain.lab.sim.motion.workspace.configs import ( SamplingConfig, VisualizationConfig, DimensionConstraint, @@ -88,7 +88,7 @@ def setup_method(self): self.setup_simulation() def test_joint_workspace_analyze(self): - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( WorkspaceAnalyzer, WorkspaceAnalyzerConfig, AnalysisMode, @@ -119,7 +119,7 @@ def setup_method(self): self.setup_simulation() def test_cartesian_workspace_analyze(self): - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( WorkspaceAnalyzer, WorkspaceAnalyzerConfig, AnalysisMode, @@ -168,7 +168,7 @@ def setup_method(self): self.setup_simulation() def test_plane_workspace_analyze(self): - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( WorkspaceAnalyzer, WorkspaceAnalyzerConfig, AnalysisMode, @@ -223,7 +223,7 @@ def setup_method(self): def test_all_modes_consistency(self): """Test basic functionality and result consistency of all three modes""" - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( WorkspaceAnalyzer, WorkspaceAnalyzerConfig, AnalysisMode, diff --git a/tests/sim/workspace/test_cache.py b/tests/sim/motion/workspace/test_cache.py similarity index 97% rename from tests/sim/workspace/test_cache.py rename to tests/sim/motion/workspace/test_cache.py index 55010e2d4..8a6d636ee 100644 --- a/tests/sim/workspace/test_cache.py +++ b/tests/sim/motion/workspace/test_cache.py @@ -22,7 +22,7 @@ import pytest import torch -from embodichain.lab.sim.workspace.caches.results_cache import ( +from embodichain.lab.sim.motion.workspace.caches.results_cache import ( ResultsCache, compute_cache_key, deserialize_results, @@ -196,7 +196,7 @@ def test_compute_cache_key_stability_and_sensitivity(): def _write_cache_entry(tmp_path, results, mode): - from embodichain.lab.sim.workspace.caches import ( + from embodichain.lab.sim.motion.workspace.caches import ( ResultsCache, compute_cache_key, ) @@ -402,7 +402,7 @@ def visualize(self, **kwargs): def test_point_cloud_visualizer_publishes_persistent_viser_overlay(): """Workspace colors and points are published through SimulationManager.""" - from embodichain.lab.sim.workspace.visualizers import ( + from embodichain.lab.sim.motion.workspace.visualizers import ( PointCloudVisualizer, ) @@ -435,7 +435,7 @@ def set_visualization_overlays(self, overlays): def test_point_cloud_visualizer_uses_sim_manager_point_cloud_api(): """The native workspace backend delegates point-cloud creation to the manager.""" - from embodichain.lab.sim.workspace.visualizers import ( + from embodichain.lab.sim.motion.workspace.visualizers import ( PointCloudVisualizer, ) @@ -471,7 +471,7 @@ def test_workspace_analyzer_prefers_configured_viser_backend(): """Automatic workspace visualization selects Viser before local backends.""" from types import SimpleNamespace - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( WorkspaceAnalyzer, ) @@ -649,7 +649,7 @@ def test_build_preset_robot_cfg_invalid_control_part(): def test_build_analyzer_config_cartesian_with_bounds(): """Cartesian config wires bounds, sampler, ik_samples and cache_dir.""" from embodichain.lab.scripts.analyze_workspace import build_analyzer_config - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( AnalysisMode, ) @@ -667,7 +667,7 @@ def test_build_analyzer_config_cartesian_with_bounds(): def test_build_analyzer_config_plane_sets_plane_params(): """Plane mode populates plane_normal / plane_point tensors.""" from embodichain.lab.scripts.analyze_workspace import build_analyzer_config - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( AnalysisMode, ) @@ -742,11 +742,11 @@ def teardown_method(self): def test_cache_hit_skips_recompute(self, tmp_path): """A second analyze() with force_recompute=False loads from cache.""" - from embodichain.lab.sim.workspace.configs import ( + from embodichain.lab.sim.motion.workspace.configs import ( SamplingConfig, VisualizationConfig, ) - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( AnalysisMode, WorkspaceAnalyzer, WorkspaceAnalyzerConfig, @@ -787,11 +787,11 @@ def _boom(*_args, **_kwargs): def test_force_recompute_ignores_cache(self, tmp_path): """force_recompute=True recomputes even when a cache entry exists.""" - from embodichain.lab.sim.workspace.configs import ( + from embodichain.lab.sim.motion.workspace.configs import ( SamplingConfig, VisualizationConfig, ) - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( AnalysisMode, WorkspaceAnalyzer, WorkspaceAnalyzerConfig, @@ -825,11 +825,11 @@ def _spy(*args, **kwargs): def test_different_params_use_different_cache_keys(self, tmp_path): """Changing num_samples produces a different cache entry.""" - from embodichain.lab.sim.workspace.configs import ( + from embodichain.lab.sim.motion.workspace.configs import ( SamplingConfig, VisualizationConfig, ) - from embodichain.lab.sim.workspace.analyzer import ( + from embodichain.lab.sim.motion.workspace.analyzer import ( AnalysisMode, WorkspaceAnalyzer, WorkspaceAnalyzerConfig, @@ -858,7 +858,7 @@ def run(n): def _cache_cfg(tmp_path): - from embodichain.lab.sim.workspace.configs import CacheConfig + from embodichain.lab.sim.motion.workspace.configs import CacheConfig return CacheConfig(enabled=True, cache_dir=str(tmp_path / "results")) diff --git a/tests/sim/workspace/test_runtime.py b/tests/sim/motion/workspace/test_runtime.py similarity index 98% rename from tests/sim/workspace/test_runtime.py rename to tests/sim/motion/workspace/test_runtime.py index 26c158db1..35c1d3829 100644 --- a/tests/sim/workspace/test_runtime.py +++ b/tests/sim/motion/workspace/test_runtime.py @@ -24,7 +24,7 @@ import torch from embodichain.lab.sim.objects import Robot -from embodichain.lab.sim.workspace import RobotWorkspace, RobotWorkspaceCfg +from embodichain.lab.sim.motion.workspace import RobotWorkspace, RobotWorkspaceCfg def test_load_joint_workspace_cache(tmp_path): diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index 0471f8119..7a16984b4 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -44,6 +44,17 @@ NUM_ARENAS = 10 +@pytest.mark.no_sim +def test_root_link_name_reads_registered_native_entity() -> None: + """Initial-state capture resolves the actual root through native entities.""" + articulation = object.__new__(Articulation) + articulation._entities = [ + SimpleNamespace(get_root_link_name=lambda: "fixed_robot_base") + ] + + assert articulation.root_link_name == "fixed_robot_base" + + class _GravityEntity: """Record native gravity calls for an articulation test double.""" diff --git a/tests/sim/objects/test_dual_arm.py b/tests/sim/objects/test_dual_arm.py index ae32f02c9..67745bced 100644 --- a/tests/sim/objects/test_dual_arm.py +++ b/tests/sim/objects/test_dual_arm.py @@ -27,7 +27,7 @@ resolve_mounts, ) from embodichain.lab.sim.robots.ur_robot import URRobotCfg -from embodichain.lab.sim.solvers import URSolverCfg +from embodichain.lab.sim.motion.solvers import URSolverCfg # --------------------------------------------------------------------------- # # resolve_mounts diff --git a/tests/sim/objects/test_gizmo.py b/tests/sim/objects/test_gizmo.py index 2a525f025..d15ea5e13 100644 --- a/tests/sim/objects/test_gizmo.py +++ b/tests/sim/objects/test_gizmo.py @@ -418,7 +418,7 @@ def __init__(self, urdf_path: Path, backend: str) -> None: self.root_pose[1:3, 3] = [-0.2, 0.6] if backend == "embodichain": pytest.importorskip("pink") - from embodichain.lab.sim.solvers import PinkSolverCfg + from embodichain.lab.sim.motion.solvers import PinkSolverCfg self.solver = PinkSolverCfg( urdf_path=str(urdf_path), diff --git a/tests/sim/objects/test_robot_cfg.py b/tests/sim/objects/test_robot_cfg.py index ca025a119..6a81f9873 100644 --- a/tests/sim/objects/test_robot_cfg.py +++ b/tests/sim/objects/test_robot_cfg.py @@ -24,7 +24,7 @@ JointDrivePropertiesCfg, RobotCfg, ) -from embodichain.lab.sim.workspace import RobotWorkspaceCfg +from embodichain.lab.sim.motion.workspace import RobotWorkspaceCfg from embodichain.lab.sim.robots.dexforce_w1 import DexforceW1Cfg from embodichain.lab.sim.robots.dexforce_w1.params import W1ArmKineParams from embodichain.lab.sim.robots.dexforce_w1.types import ( @@ -43,7 +43,7 @@ build_dexforce_w1_assembly_urdf_cfg, build_dexforce_w1_control_parts, ) -from embodichain.lab.sim.solvers import SRSSolverCfg +from embodichain.lab.sim.motion.solvers import SRSSolverCfg from embodichain.utils import configclass from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg @@ -430,7 +430,7 @@ def test_robotcfg_to_dict_roundtrip(): from embodichain.lab.sim.robots.cobotmagic import CobotMagicCfg -from embodichain.lab.sim.solvers import OPWSolverCfg +from embodichain.lab.sim.motion.solvers import OPWSolverCfg def test_cobotmagic_from_dict_and_roundtrip(): @@ -524,7 +524,7 @@ def test_cobotmagic_pk_dof_matches_control_parts(): # --------------------------------------------------------------------------- # from embodichain.lab.sim.robots.ur_robot import URRobotCfg -from embodichain.lab.sim.solvers import URSolverCfg +from embodichain.lab.sim.motion.solvers import URSolverCfg UR_TYPES = ["ur3", "ur3e", "ur5", "ur5e", "ur10", "ur10e"] diff --git a/tests/toolkits/test_grasp_pose_generator.py b/tests/toolkits/test_grasp_pose_generator.py index ed09a6714..11800466d 100644 --- a/tests/toolkits/test_grasp_pose_generator.py +++ b/tests/toolkits/test_grasp_pose_generator.py @@ -32,7 +32,7 @@ from embodichain.lab.sim.utility.action_utils import interpolate_with_distance from embodichain.toolkits.graspkit import ParallelJawGripperModelCfg from embodichain.lab.sim.shapes import MeshCfg -from embodichain.lab.sim.solvers import PytorchSolverCfg +from embodichain.lab.sim.motion.solvers import PytorchSolverCfg from embodichain.data import get_data_path from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.utils import logger From 92c9bd8a1659e3130aaf37c0c96b85640e92a757 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 6 Sep 2026 17:04:22 +0000 Subject: [PATCH 2/2] fix(examples): configure trajectory visualization --- .../sim/motion/trajectory_generation/cube_grasp_parallel.py | 2 ++ .../sim/motion/trajectory_generation/cube_pickup_collection.py | 2 ++ examples/sim/motion/trajectory_generation/free_motion.py | 2 ++ 3 files changed, 6 insertions(+) diff --git a/examples/sim/motion/trajectory_generation/cube_grasp_parallel.py b/examples/sim/motion/trajectory_generation/cube_grasp_parallel.py index 12dce2b4f..1b722a048 100644 --- a/examples/sim/motion/trajectory_generation/cube_grasp_parallel.py +++ b/examples/sim/motion/trajectory_generation/cube_grasp_parallel.py @@ -69,6 +69,7 @@ from embodichain.lab.sim.objects import RigidObjectCfg from embodichain.lab.sim.sensors import CameraCfg from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.visualization import VisualizationCfg from embodichain.utils.math import look_at_to_pose from scripts.tutorials.atomic_action.tutorial_utils import ( create_antipodal_semantics, @@ -321,6 +322,7 @@ def run_cube_grasp_parallel( physics_dt=_PHYSICS_DT, arena_space=4.0, gpu_id=cuda_device, + visualization=VisualizationCfg(), ) ) generator = writer = None diff --git a/examples/sim/motion/trajectory_generation/cube_pickup_collection.py b/examples/sim/motion/trajectory_generation/cube_pickup_collection.py index 3a9084a3f..652e6f53e 100644 --- a/examples/sim/motion/trajectory_generation/cube_pickup_collection.py +++ b/examples/sim/motion/trajectory_generation/cube_pickup_collection.py @@ -49,6 +49,7 @@ ValidationCheck, ValidationResult, ) +from embodichain.lab.visualization import VisualizationCfg from embodichain.lab.trajectory_generation.execution import QposRolloutExecutor from embodichain.lab.trajectory_generation.initial_state import ( FixedSceneHost, @@ -140,6 +141,7 @@ def run_cube_pickup_collection( physics_dt=_PHYSICS_DT, arena_space=4.0, gpu_id=cuda_device, + visualization=VisualizationCfg(), ) ) generator = host = sink = writer = None diff --git a/examples/sim/motion/trajectory_generation/free_motion.py b/examples/sim/motion/trajectory_generation/free_motion.py index 98dde54b9..eba65403d 100644 --- a/examples/sim/motion/trajectory_generation/free_motion.py +++ b/examples/sim/motion/trajectory_generation/free_motion.py @@ -67,6 +67,7 @@ from embodichain.lab.sim.robots import FrankaPandaCfg, URRobotCfg from embodichain.lab.sim.shapes import CubeCfg from embodichain.lab.sim.sensors import CameraCfg +from embodichain.lab.visualization import VisualizationCfg from embodichain.lab.trajectory_generation.execution import QposRolloutExecutor from embodichain.lab.trajectory_generation.initial_state import ( FixedSceneHost, @@ -128,6 +129,7 @@ def run_free_motion( num_envs=1, physics_dt=physics_dt, gpu_id=cuda_device, + visualization=VisualizationCfg(), ) ) host = sink = generator = video_writer = None