diff --git a/.agents/skills/add-solver/SKILL.md b/.agents/skills/add-solver/SKILL.md index eabad3e25..d6d2a1051 100644 --- a/.agents/skills/add-solver/SKILL.md +++ b/.agents/skills/add-solver/SKILL.md @@ -37,6 +37,10 @@ Plus two registration edits: - Add the docs page to the toctree in `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 ### 1. Gather Solver Requirements diff --git a/.agents/skills/review-pr/references/review-matrix.md b/.agents/skills/review-pr/references/review-matrix.md index 3627b79ae..5a69854a2 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,expansion}/**`, +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 689fd82c7..363a83cbd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -29,6 +29,7 @@ bundles official tasks as the **`embodichain_tasks`** import package. | Simulation world, objects, sensors, solvers, planning | `embodichain/lab/sim/` | | Gym environments and manager functors | `embodichain/lab/gym/` | | Task Program language, semantics, compiler, runtime and integrations | `embodichain/lab/task_program/` | +| Fixed-scene trajectory host, rollout and persistence | `embodichain/lab/trajectory_generation/` | | Browser visualization | `embodichain/lab/visualization/` | | RL algorithms, policies, collectors and trainers | `embodichain/learning/rl/` | | Real-device controllers / standalone tools | `embodichain/lab/devices/`, `embodichain/toolkits/` | diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index b670e0448..6342c2558 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -21,6 +21,12 @@ topics: - embodichain/lab/sim/objects/articulation.py - embodichain/lab/sim/objects/gizmo.py - embodichain/lab/gym/envs/base_env.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/expansion/__init__.py + - embodichain/lab/trajectory_generation/initial_state.py + - embodichain/lab/trajectory_generation/integrations/sim.py watch_paths: [embodichain/lab/sim/sim_manager.py, embodichain/lab/sim/cfg.py, embodichain/lab/sim/objects/, tests/sim/] related_topics: [env-framework, robot-system, sensor-system, sim-visualization, ik-solvers, motion-planning, @@ -28,9 +34,34 @@ topics: status: active - id: env-framework title: Environment Framework - aliases: [env framework, environment framework, task registration, list task, 环境框架, 任务环境] - keywords: [BaseEnv, EmbodiedEnv, EnvCfg, register_env, list-task, run-env, target_control_frequency, - sim_steps_per_control, step_dt, environment.component, embodiment.component, ControllerAction, EnvProfiler] + aliases: + - env framework + - environment framework + - task registration + - list task + - 环境框架 + - 任务环境 + - generation lease + - controlled episode preparation + - 受控初态准备 + keywords: + - BaseEnv + - EmbodiedEnv + - EnvCfg + - register_env + - list-task + - run-env + - target_control_frequency + - sim_steps_per_control + - step_dt + - environment.component + - embodiment.component + - ControllerAction + - EnvProfiler + - acquire_generation_lease + - prepare_generation_episode + - generation_epoch + - FixedSceneHost paths: [topics/env-framework/env-framework.md] source_of_truth: - embodichain/cli/main.py @@ -43,6 +74,7 @@ topics: - embodichain/lab/gym/envs/embodied_env.py - embodichain/lab/gym/envs/demo.py - embodichain/lab/gym/utils/profiler.py + - embodichain/lab/trajectory_generation/initial_state.py watch_paths: [embodichain/lab/gym/envs/base_env.py, embodichain/lab/gym/envs/embodied_env.py, embodichain/lab/gym/envs/demo.py, embodichain/lab/gym/utils/, embodichain/lab/gym/envs/task_program/, embodichain/cli/, embodichain/lab/scripts/run_env.py, embodichain_tasks/, tests/gym/envs/] @@ -83,20 +115,51 @@ topics: - embodichain/lab/sim/motion/solvers/null_space_posture_task.py - embodichain/lab/sim/utility/solver_utils.py - embodichain/compute/kinematics/ + - embodichain/lab/sim/cfg.py + - embodichain/lab/sim/motion/__init__.py + - embodichain/lab/sim/motion/solvers/srs_solver.py + - embodichain/lab/sim/motion/solvers/opw_solver.py + - embodichain/lab/sim/motion/solvers/pinocchio_solver.py + - embodichain/lab/sim/motion/solvers/differential_solver.py watch_paths: [embodichain/lab/sim/motion/solvers/, embodichain/lab/sim/utility/solver_utils.py, embodichain/compute/kinematics/, tests/sim/motion/solvers/, tests/compute/] related_topics: [simulation-system, robot-system, motion-planning] status: active - id: robot-system title: Robot System - aliases: [robot system, robot config, 机器人配置, 机器人系统] - keywords: [RobotCfg, control_parts, build_pk_serial_chain, DexforceW1Cfg, CobotMagicCfg, FrankaPandaCfg, - URRobotCfg, DualArmRobotCfg, merge_robot_cfg, JointDrivePropertiesCfg] + aliases: + - robot system + - robot config + - 机器人配置 + - 机器人系统 + - robot workspace + - workspace analysis + - 机器人工作空间 + keywords: + - RobotCfg + - control_parts + - build_pk_serial_chain + - DexforceW1Cfg + - CobotMagicCfg + - FrankaPandaCfg + - URRobotCfg + - DualArmRobotCfg + - merge_robot_cfg + - JointDrivePropertiesCfg + - workspace + - RobotWorkspaceCfg + - RobotWorkspace + - WorkspaceAnalyzer + - reachability paths: [topics/robot-system/robot-system.md] source_of_truth: - 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 watch_paths: [embodichain/lab/sim/robots/, embodichain/lab/sim/objects/robot.py, embodichain/lab/sim/cfg.py] related_topics: [simulation-system, ik-solvers, motion-planning, sensor-system, sim-visualization, robot-workspace] status: active @@ -137,9 +200,26 @@ topics: - id: motion-planning title: Motion Planning aliases: [motion planning, trajectory planning, motion expansion, trajectory expansion, trajectory augmentation, fixed scene trajectory augmentation, 轨迹扩增, 运动规划, 轨迹规划] - keywords: [BasePlanner, PlanState, PlanResult, MotionGenerator, ToppraPlanner, CuroboPlanner, NeuralPlanner, - expansion, GenerationSession, TrajectoryAugmentationCfg, CandidateTrajectoryBatch, - collision world, compute trajectory, trajectory resampling, trajectory warping] + keywords: + - BasePlanner + - PlanState + - PlanResult + - MotionGenerator + - ToppraPlanner + - CuroboPlanner + - NeuralPlanner + - expansion + - GenerationSession + - TrajectoryAugmentationCfg + - CandidateTrajectoryBatch + - collision world + - compute trajectory + - trajectory resampling + - trajectory warping + - GenerationRunner + - FixedSceneHost + - QposRolloutExecutor + - trajectory_generation paths: [topics/motion-planning/motion-planning.md] source_of_truth: - embodichain/lab/sim/motion/planners/base_planner.py @@ -152,6 +232,25 @@ topics: - embodichain/lab/sim/motion/expansion/ - embodichain/compute/trajectory/ - embodichain/lab/sim/utility/action_utils.py + - embodichain/lab/sim/motion/__init__.py + - embodichain/lab/sim/motion/planners/neural_planner.py + - embodichain/lab/sim/motion/expansion/__init__.py + - embodichain/lab/sim/motion/expansion/contracts.py + - embodichain/lab/sim/motion/expansion/cfg.py + - embodichain/lab/sim/motion/expansion/operators.py + - embodichain/lab/sim/motion/expansion/coverage.py + - embodichain/lab/sim/motion/expansion/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 watch_paths: [embodichain/lab/sim/motion/motion_generator.py, tests/sim/motion/test_motion_generator.py, tests/sim/motion/test_motion_generator_batched.py, embodichain/lab/sim/motion/planners/, embodichain/lab/sim/motion/expansion/, tests/sim/motion/expansion/, embodichain/lab/sim/utility/action_utils.py, embodichain/compute/trajectory/, diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 27c728519..dce410679 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -124,3 +124,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/atomic-actions/execution.md b/agent_context/topics/atomic-actions/execution.md index ab286e432..107566fc9 100644 --- a/agent_context/topics/atomic-actions/execution.md +++ b/agent_context/topics/atomic-actions/execution.md @@ -47,6 +47,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 @@ -90,6 +100,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, diff --git a/agent_context/topics/env-framework/env-framework.md b/agent_context/topics/env-framework/env-framework.md index e9d824ab2..0b33f2c5e 100644 --- a/agent_context/topics/env-framework/env-framework.md +++ b/agent_context/topics/env-framework/env-framework.md @@ -27,6 +27,8 @@ component ownership, path resolution, config-owned IDs, and task listing. Read [execution](execution.md) for hooks, bridge acceptance, reset ordering, wrappers, and replay. Read [profiling](profiling.md) only for instrumentation. +For fixed-scene collection, see [generation preparation and explicit demo candidates](execution.md#fixed-scene-generation-preparation-base_envpy-embodied_envpy). The host owns a full-batch generation lease and validated preparation epochs. + ## Timing contract `BaseEnv._configure_timing()` resolves `EnvCfg` before constructing the scene: diff --git a/agent_context/topics/env-framework/execution.md b/agent_context/topics/env-framework/execution.md index 5aad0290a..e78a242d3 100644 --- a/agent_context/topics/env-framework/execution.md +++ b/agent_context/topics/env-framework/execution.md @@ -49,6 +49,68 @@ Read this when the request needs these details. [Topic overview](env-framework.m - 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/motion-planning/motion-planning.md b/agent_context/topics/motion-planning/motion-planning.md index 9e8853edd..6a505e438 100644 --- a/agent_context/topics/motion-planning/motion-planning.md +++ b/agent_context/topics/motion-planning/motion-planning.md @@ -60,6 +60,8 @@ the object nor certifies the grasp. Focused augmentation tests live under `tests/sim/motion/expansion/`. +Read [fixed-scene trajectory generation](trajectory-generation.md) for host restoration, qpos rollout, offline PickUp sources, contact validation, and confirmed dataset persistence. + ## Choose the owning layer - `BasePlanner` and `PlanState` / `PlanResult` define planning interfaces. diff --git a/agent_context/topics/motion-planning/trajectory-generation.md b/agent_context/topics/motion-planning/trajectory-generation.md new file mode 100644 index 000000000..c0d1df61e --- /dev/null +++ b/agent_context/topics/motion-planning/trajectory-generation.md @@ -0,0 +1,97 @@ +# Fixed-scene trajectory generation + +[Motion planning overview](motion-planning.md). + +`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/`. diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index c73d6262a..22d5da702 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -123,6 +123,21 @@ 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..a0b02b6b2 --- /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_expansion_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/expansion/`，包含 `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/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_expansion_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..3be40342e --- /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_expansion_design.md)，2026-09-05。 +- 代码核对基线：`fb7228e2`。首批实现基于该版本，已运行 CPU 行为测试；尚未完成四组合真实仿真和性能验收。 +- 交付原则：先完成可验收、可持久化的四条执行路径，再优化吞吐，最后扩大运动覆盖。 + +## 当前实施进度（2026-09-05） + +| 计划项 | 已实现 | 尚待实现与验收 | +|---|---|---| +| PR 1 | `lab.sim.motion.expansion`：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` 与 `expansion` 统一归入 `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/expansion/`；验证算法无直接 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/expansion/ + 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/expansion/`：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.motion.expansion.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.expansion.rst index 187e126b2..ee41c8bec 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.expansion.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.expansion.rst @@ -8,11 +8,13 @@ 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. Physical -initial-state restoration, planning, rollout execution, task validation, and -episode persistence must be supplied by separate host integrations. Those -integrations provide actual observations and commands, validation evidence, -and persistence confirmations; this package does not instantiate them. +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.expansion diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst index 92a22519e..4e1244155 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst @@ -1,4 +1,5 @@ embodichain.lab.sim.motion.planners +=================================== ========================================== .. automodule:: embodichain.lab.sim.motion.planners diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst index bc70ae02e..f91e75214 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst @@ -1,4 +1,5 @@ embodichain.lab.sim.motion.solvers +================================== ========================================== .. automodule:: embodichain.lab.sim.motion.solvers 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.trajectory_generation.rst b/docs/source/api_reference/embodichain/embodichain.lab.trajectory_generation.rst new file mode 100644 index 000000000..abcff3f74 --- /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.expansion`. +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 253938d7a..3e8aa8dc8 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -612,21 +612,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,94 +814,20 @@ 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.motion.planners.base_planner ------------------------------------------------ @@ -954,6 +865,17 @@ embodichain.lab.sim.motion.planners.curobo.curobo_yaml generate_curobo_robot_yaml generate_curobo_world_yaml +embodichain.lab.sim.motion.motion_generator +---------------------------------------------------- + +.. currentmodule:: embodichain.lab.sim.motion.motion_generator + +.. autosummary:: + + MotionGenerator + MotionGenCfg + MotionGenOptions + embodichain.lab.sim.motion.planners.neural_planner -------------------------------------------------- @@ -993,457 +915,546 @@ embodichain.lab.sim.motion.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 + HaltonSampler -.. automodule:: embodichain.lab.task_program.semantics.calls - :members: - :no-index: +embodichain.lab.sim.motion.workspace.samplers.importance_sampler +---------------------------------------------------------------- -embodichain.lab.task_program.semantics.integration ------------------------------------------------------- +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.importance_sampler -.. currentmodule:: embodichain.lab.task_program.semantics.integration +.. autosummary:: + + ImportanceSampler + +embodichain.lab.sim.motion.workspace.samplers.iniform_sampler +------------------------------------------------------------- + +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.iniform_sampler .. autosummary:: - BoundSemanticCall - LinkedSemanticCall - PathPart - SceneEntityManifest - SceneManifest - SemanticDiagnostic - SemanticIntegrationManifest - SemanticValidationError + UniformSampler -.. automodule:: embodichain.lab.task_program.semantics.integration - :members: - :no-index: +embodichain.lab.sim.motion.workspace.samplers.lhs_sampler +--------------------------------------------------------- -embodichain.lab.task_program.semantics.effects --------------------------------------------------- +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.lhs_sampler -.. automodule:: embodichain.lab.task_program.semantics.effects - :members: - :no-index: +.. autosummary:: -embodichain.lab.task_program.semantics.evidence ---------------------------------------------------- + LatinHypercubeSampler -.. automodule:: embodichain.lab.task_program.semantics.evidence - :members: - :no-index: +embodichain.lab.sim.motion.workspace.samplers.random_sampler +------------------------------------------------------------ -embodichain.lab.task_program.semantics.profiles ---------------------------------------------------- +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.random_sampler -.. currentmodule:: embodichain.lab.task_program.semantics.profiles +.. autosummary:: + + RandomSampler + +embodichain.lab.sim.motion.workspace.samplers.sobol_sampler +----------------------------------------------------------- + +.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.sobol_sampler .. autosummary:: - AmbiguousSkillBindingError - BoundRobotSkillProfile - ControlPartEndpoint - ControlPartEndpointAdapter - EffectAssurance - EndpointResolution - ProfileValidationError - ResourceEndpoint - ResourceEndpointAdapter - ResolvedRobotResource - ResolvedResourceEndpoint - ResolvedSkillBinding - ResourceBinding - ResourceClaim - RobotResource - RobotSkillProfile - SkillPolicyPreset - UnsupportedSkillError - WorkflowRecoveryPolicy + SobolSampler -embodichain.lab.task_program.semantics.scene ------------------------------------------------- +embodichain.lab.sim.motion.workspace.visualizers.axis_visualizer +---------------------------------------------------------------- -.. currentmodule:: embodichain.lab.task_program.semantics.scene +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.axis_visualizer .. 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 + AxisVisualizer -.. automodule:: embodichain.lab.task_program.semantics.scene - :members: - :no-index: +embodichain.lab.sim.motion.workspace.visualizers.base_visualizer +---------------------------------------------------------------- -embodichain.lab.sim.motion.solvers.neural_ik_solver ---------------------------------------------------- +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.base_visualizer -.. currentmodule:: embodichain.lab.sim.motion.solvers.neural_ik_solver +.. autosummary:: + + IVisualizer + BaseVisualizer + +embodichain.lab.sim.motion.workspace.visualizers.point_cloud_visualizer +----------------------------------------------------------------------- + +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.point_cloud_visualizer .. autosummary:: - NeuralIKSolverCfg - NeuralIKSolver + PointCloudVisualizer -embodichain.lab.sim.motion.solvers.null_space_posture_task ----------------------------------------------------------- +embodichain.lab.sim.motion.workspace.visualizers.sphere_visualizer +------------------------------------------------------------------ -.. currentmodule:: embodichain.lab.sim.motion.solvers.null_space_posture_task +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.sphere_visualizer .. autosummary:: - NullSpacePostureTask + SphereVisualizer -embodichain.lab.sim.motion.solvers.pink_solver ----------------------------------------------- +embodichain.lab.sim.motion.workspace.visualizers.visualizer_factory +------------------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.solvers.pink_solver +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.visualizer_factory .. autosummary:: - PinkSolver - PinkSolverCfg + VisualizerFactory + create_visualizer -embodichain.lab.sim.motion.solvers.srs_solver ---------------------------------------------- +embodichain.lab.sim.motion.workspace.visualizers.voxel_visualizer +----------------------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.solvers.srs_solver +.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.voxel_visualizer .. autosummary:: - SRSSolver - SRSSolverCfg + VoxelVisualizer -embodichain.lab.sim.utility.render_utils +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.motion.workspace.caches.cache_utils -------------------------------------------------------- +embodichain.lab.sim.objects.cloth_object +---------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.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.motion.workspace.caches.results_cache ---------------------------------------------------------- +embodichain.lab.sim.objects.constraint +-------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.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.motion.workspace.constraints.base_constraint ----------------------------------------------------------------- +embodichain.lab.sim.objects.gizmo +--------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.constraints.base_constraint +.. 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:: - IConstraintChecker - BaseConstraintChecker + Gizmo + GizmoCfg + create_robot_ik_gizmo_controller -embodichain.lab.sim.motion.workspace.constraints.workspace_constraint ---------------------------------------------------------------------- +embodichain.lab.sim.objects.rigid_object +---------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.constraints.workspace_constraint +.. currentmodule:: embodichain.lab.sim.objects.rigid_object .. autosummary:: - WorkspaceConstraintChecker + RigidBodyData + RigidObject + RigidObjectCfg -embodichain.lab.sim.motion.workspace.samplers.base_sampler ----------------------------------------------------------- +embodichain.lab.sim.objects.rigid_object_group +---------------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.base_sampler +.. currentmodule:: embodichain.lab.sim.objects.rigid_object_group .. autosummary:: - ISampler - BaseSampler + RigidBodyGroupData + RigidObjectGroup + RigidObjectGroupCfg -embodichain.lab.sim.motion.workspace.samplers.gaussian_sampler --------------------------------------------------------------- +embodichain.lab.sim.objects.robot +--------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.gaussian_sampler +.. currentmodule:: embodichain.lab.sim.objects.robot .. autosummary:: - GaussianSampler + ControlGroup + Robot -embodichain.lab.sim.motion.workspace.samplers.halton_sampler ------------------------------------------------------------- +embodichain.lab.sim.objects.soft_object +--------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.halton_sampler +.. currentmodule:: embodichain.lab.sim.objects.soft_object .. autosummary:: - HaltonSampler + SoftBodyData + SoftObject + SoftObjectCfg -embodichain.lab.sim.motion.workspace.samplers.importance_sampler ----------------------------------------------------------------- +embodichain.lab.sim.robots.cobotmagic +------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.importance_sampler +.. currentmodule:: embodichain.lab.sim.robots.cobotmagic .. autosummary:: - ImportanceSampler + CobotMagicCfg -embodichain.lab.sim.motion.workspace.samplers.iniform_sampler -------------------------------------------------------------- +embodichain.lab.sim.robots.dexforce_w1.hand_specs +------------------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.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.motion.workspace.samplers.lhs_sampler ---------------------------------------------------------- +embodichain.lab.sim.robots.dexforce_w1.specs +-------------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.lhs_sampler +.. currentmodule:: embodichain.lab.sim.robots.dexforce_w1.specs .. autosummary:: - LatinHypercubeSampler + W1VersionSpec + get_w1_version_spec -embodichain.lab.sim.motion.workspace.samplers.random_sampler ------------------------------------------------------------- +embodichain.lab.sim.robots.dexforce_w1.types +-------------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.samplers.random_sampler +.. currentmodule:: embodichain.lab.sim.robots.dexforce_w1.types .. autosummary:: - RandomSampler + DexforceW1Version + DexforceW1HandVersion + DexforceW1ArmSide + DexforceW1Type + DexforceW1HandBrand -embodichain.lab.sim.motion.workspace.samplers.sobol_sampler ------------------------------------------------------------ +embodichain.lab.sim.robots.dexforce_w1.utils +-------------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.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.motion.workspace.visualizers.axis_visualizer ----------------------------------------------------------------- +embodichain.lab.sim.robots.dual_arm +----------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.axis_visualizer +.. currentmodule:: embodichain.lab.sim.robots.dual_arm .. autosummary:: - AxisVisualizer + DualArmRobotCfg + build_dual_arm_cfg + resolve_mounts -embodichain.lab.sim.motion.workspace.visualizers.base_visualizer ----------------------------------------------------------------- +embodichain.lab.sim.robots.franka_panda +--------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.base_visualizer +.. currentmodule:: embodichain.lab.sim.robots.franka_panda .. autosummary:: - IVisualizer - BaseVisualizer + FrankaPandaCfg -embodichain.lab.sim.motion.workspace.visualizers.point_cloud_visualizer ------------------------------------------------------------------------ +embodichain.lab.sim.robots.ur_robot +----------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.point_cloud_visualizer +.. currentmodule:: embodichain.lab.sim.robots.ur_robot .. autosummary:: - PointCloudVisualizer + URRobotCfg -embodichain.lab.sim.motion.workspace.visualizers.sphere_visualizer ------------------------------------------------------------------- +embodichain.lab.sim.sensors.camera +---------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.sphere_visualizer +.. currentmodule:: embodichain.lab.sim.sensors.camera .. autosummary:: - SphereVisualizer + Camera + CameraCfg -embodichain.lab.sim.motion.workspace.visualizers.visualizer_factory -------------------------------------------------------------------- +embodichain.lab.sim.sim_manager +------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.visualizer_factory +.. currentmodule:: embodichain.lab.sim.sim_manager .. autosummary:: - VisualizerFactory - create_visualizer + SIM_CACHE_DIR + MATERIAL_CACHE_DIR + CONVEX_DECOMP_DIR + REACHABLE_XPOS_DIR -embodichain.lab.sim.motion.workspace.visualizers.voxel_visualizer ------------------------------------------------------------------ +embodichain.lab.sim.utility.render_utils +---------------------------------------- -.. currentmodule:: embodichain.lab.sim.motion.workspace.visualizers.voxel_visualizer +.. currentmodule:: embodichain.lab.sim.utility.render_utils .. autosummary:: - VoxelVisualizer + select_default_renderer + +embodichain.lab.task_program.semantics.calls +------------------------------------------------ + +.. currentmodule:: embodichain.lab.task_program.semantics.calls + +.. 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: + +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.task_program.semantics.integration + +.. autosummary:: + + 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/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/motion/index.rst b/docs/source/overview/sim/motion/index.rst index 2d8a22438..6a5785790 100644 --- a/docs/source/overview/sim/motion/index.rst +++ b/docs/source/overview/sim/motion/index.rst @@ -27,12 +27,12 @@ public imports still use the normal ``lab`` and ``sim`` initialization path. - Explicit qpos templates and candidates, constrained geometric and timing variation, measured coverage, and bounded generation accounting. -Trajectory augmentation provides the core contracts and operators. Physical -initial-state restoration, planning, rollout execution, task validation, and -dataset persistence must be supplied by host integrations. The core accepts -their evidence and persistence receipts without creating those services. See the +Trajectory augmentation provides the core contracts and operators. The +:doc:`generation host layer ` supplies +full-batch restoration, qpos execution, measured validation and confirmed +persistence for free motion and offline atomic PickUp sources. See the :doc:`augmentation API ` -for the implemented boundaries. +for the core boundaries. Migrating Existing Imports -------------------------- @@ -55,7 +55,8 @@ imports, string-based module references in configuration, and custom extensions: Their nested modules follow the same mapping. Solver ``class_type`` names such as ``URSolver`` still resolve through ``RobotCfg.from_dict()``. Existing tests and examples now live under ``tests/sim/motion/`` and ``examples/sim/motion/``. -Warp kinematics kernels remain at ``embodichain.utils.warp.kinematics``. +Warp kinematics kernels live under ``embodichain.compute.kinematics._warp``; +``embodichain.utils.warp.kinematics`` retains compatibility exports. .. toctree:: :maxdepth: 1 diff --git a/docs/source/overview/sim/motion/planners/curobo_planner.md b/docs/source/overview/sim/motion/planners/curobo_planner.md index eefd3066a..6133b1884 100644 --- a/docs/source/overview/sim/motion/planners/curobo_planner.md +++ b/docs/source/overview/sim/motion/planners/curobo_planner.md @@ -292,13 +292,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 diff --git a/docs/source/overview/trajectory_generation.rst b/docs/source/overview/trajectory_generation.rst new file mode 100644 index 000000000..f0f94e5a6 --- /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.expansion` 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/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..8c8fa83b0 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.expansion 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.expansion 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/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index ef03d926d..1bec94d8f 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 @@ -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 a6a54fb06..3107a1b48 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 @@ -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/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/motion/planners/curobo/curobo_planner.py b/embodichain/lab/sim/motion/planners/curobo/curobo_planner.py index c9869d84b..750fcf710 100644 --- a/embodichain/lab/sim/motion/planners/curobo/curobo_planner.py +++ b/embodichain/lab/sim/motion/planners/curobo/curobo_planner.py @@ -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" @@ -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/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/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/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..2632571d4 --- /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.expansion 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..d1659d219 --- /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.expansion 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..7ee7c69d0 --- /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.expansion 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..b75b746cb --- /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.expansion 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.expansion.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..71900d9a9 --- /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.expansion import ( + CandidateIdentity, + CandidateTrajectoryBatch, + MotionSnapshot, + TrajectoryPhase, + TrajectoryTemplate, + ValidationCheck, + ValidationResult, +) + +if TYPE_CHECKING: + from embodichain.lab.sim.motion.motion_generator 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..92a42a126 --- /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.expansion 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..45131ee87 --- /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.expansion 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..0db0623fd --- /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.expansion 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/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..1b722a048 --- /dev/null +++ b/examples/sim/motion/trajectory_generation/cube_grasp_parallel.py @@ -0,0 +1,546 @@ +# ---------------------------------------------------------------------------- +# 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.expansion 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.lab.visualization import VisualizationCfg +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, + visualization=VisualizationCfg(), + ) + ) + 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..652e6f53e --- /dev/null +++ b/examples/sim/motion/trajectory_generation/cube_pickup_collection.py @@ -0,0 +1,418 @@ +# ---------------------------------------------------------------------------- +# 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.expansion import ( + SceneCase, + TrajectoryGenerationJobCfg, + 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, + 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, + visualization=VisualizationCfg(), + ) + ) + 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..eba65403d --- /dev/null +++ b/examples/sim/motion/trajectory_generation/free_motion.py @@ -0,0 +1,443 @@ +# ---------------------------------------------------------------------------- +# 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.motion_generator import MotionGenCfg, MotionGenerator +from embodichain.lab.sim.motion.planners import CuroboPlannerCfg, CuroboWorldCfg +from embodichain.lab.sim.motion.expansion 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.visualization import VisualizationCfg +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, + visualization=VisualizationCfg(), + ) + ) + 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/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/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..ff9db6229 --- /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.expansion 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/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..97f263c59 --- /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.expansion 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..eb959eed6 --- /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.expansion 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..2ccb51140 --- /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.expansion 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..479a968d4 --- /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.expansion 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..3afdc1ea6 --- /dev/null +++ b/tests/lab/trajectory_generation/test_planning.py @@ -0,0 +1,571 @@ +# ---------------------------------------------------------------------------- +# 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.motion_generator import MotionGenerator +from embodichain.lab.sim.motion.planners import CollisionWorldInfo +from embodichain.lab.sim.motion.expansion 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.motion_generator 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..b319a2043 --- /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.expansion 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..dc1918c84 --- /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.expansion 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..a1e2b98d0 --- /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.expansion 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 e96c35353..63dfe486b 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -2064,6 +2064,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_engine_per_env.py b/tests/sim/atomic_actions/test_engine_per_env.py index 238a9fbb2..7f5cb6041 100644 --- a/tests/sim/atomic_actions/test_engine_per_env.py +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -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/motion/planners/test_curobo_planner.py b/tests/sim/motion/planners/test_curobo_planner.py index 4609677f1..e0d6bc7f7 100644 --- a/tests/sim/motion/planners/test_curobo_planner.py +++ b/tests/sim/motion/planners/test_curobo_planner.py @@ -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", @@ -1097,3 +1116,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/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."""