diff --git a/.agents/skills/add-atomic-action/SKILL.md b/.agents/skills/add-atomic-action/SKILL.md index 02241fd46..f45823433 100644 --- a/.agents/skills/add-atomic-action/SKILL.md +++ b/.agents/skills/add-atomic-action/SKILL.md @@ -230,6 +230,16 @@ Follow these invariants: ## 4. Emit generic runtime commands when needed +For a joint-backed Skill that must also consume external or augmented motion, +declare a versioned `SkillTrajectoryContract` and implement +`_trajectory_semantics(request, context, motion) -> SkillPlanSemantics`. Call the +framework-owned `assemble_trajectory_plan()` from ordinary `_plan()` after +generating a `SkillTrajectoryInput`; external callers use the engine's +`build_plan_from_trajectory()`. Reuse the same semantic hook instead of adding +recording-specific methods. Keep file decoding/calibration in source adapters +and expansion permissions in the motion-expansion layer. See PickUp and +MoveJoints for typed geometry, event landmarks and goal checks. + Use `build_command_plan()` when a skill targets a mobile base, whole-body controller, tool, or another non-joint transport. Build immutable endpoint commands; keep live controller and device handles in the transport: diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index b322d2c84..f2b1e8c01 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -188,7 +188,7 @@ topics: via_points, nullspace_residual, perturb_approach_direction, sample_approach_cone, TIMING_PROFILES, trajectory variant, 轨迹变体, 固定路点, collision world, compute trajectory, trajectory resampling, trajectory warping, retime_to_control_grid, - JointTrajectoryPlaybackCfg, play_joint_trajectory, quaternion, xyzw, wxyz] + JointTrajectoryPlaybackCfg, play_joint_trajectory, Real2SimSourceAdapter, AnchoredTimeMap, quaternion, xyzw, wxyz] paths: [topics/motion-planning/motion-planning.md] source_of_truth: - embodichain/lab/sim/motion/planners/base_planner.py @@ -313,13 +313,14 @@ topics: - id: atomic-actions title: Atomic Skills aliases: [atomic action, atomic skill, 原子动作, 原子技能] - keywords: [AtomicActionEngine, run-atomic-action, PlanningContext, ExecutionSession, ExecutionRunner, ActionInvocation, + keywords: [AtomicActionEngine, NominalPlanProvider, SkillTrajectoryInput, SkillTrajectoryContract, SkillPlanSemantics, run-atomic-action, PlanningContext, ExecutionSession, ExecutionRunner, ActionInvocation, RecoveryPolicy, HeldObjectGuardRequest, PhaseEffectGateRequest, articulation_geometry, control-grid retiming, AffordanceSamplingContext, AffordancePoseCandidates, AffordanceSample, affordance sampling] paths: [topics/atomic-actions/atomic-actions.md] source_of_truth: - embodichain/lab/sim/atomic_actions/engine.py - embodichain/lab/sim/atomic_actions/core.py + - embodichain/lab/sim/atomic_actions/trajectory_assembly.py - embodichain/lab/sim/atomic_actions/goals.py - embodichain/lab/sim/atomic_actions/execution.py - embodichain/lab/sim/atomic_actions/verification.py @@ -339,7 +340,8 @@ topics: keywords: [TaskProgramCfg, program.yaml, integration.yaml, scene_binding, SceneManifest, SkillPolicyPreset, EffectAssurance, CompiledTaskProgram, TaskProgramDemoBridge, TaskProgramPlanTransformFactory, CombinedTaskProgramExpansionFactory, expansion_profile, execution_policy, SemanticPose, quaternion_xyzw, - joint_command_mode] + joint_command_mode, TaskProgramNominalPlanFactory, ExpertSourceProviderFactory, ExpertSourceRegistry, + Real2SimExpertSourceProviderFactory] paths: [topics/task-programs/task-programs.md] source_of_truth: - embodichain/lab/task_program/ @@ -393,7 +395,7 @@ topics: keywords: [OnlineDataEngine, OnlineDataset, sample_batch, lock_index, LeRobotRecorder, AsyncLeRobotRecorder, DepthSidecarManager, depth_meta.json, fragment_id, finalize, expert action schema, qpos_slice, qvel_slice, validate_dataset, create_split_manifest, recover_recording, episode_uuid, language_persistent, - SharedLanguageRegistry, resolve_language] + SharedLanguageRegistry, resolve_language, Real2SimTrajectory, normalize_real2sim, 真实轨迹归一化] paths: [topics/data-pipeline/data-pipeline.md] source_of_truth: - embodichain/data_pipeline/engine/data.py diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 37bcfc78f..30d2d439a 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -43,6 +43,13 @@ ActionInvocation semantic program EndpointCommandRouter + transports + clock ``` +`NominalPlanProvider` is an optional trusted complete-plan source on engine +planning/start calls. It bypasses ordinary action planning without bypassing +plan validation, and remains attached to the execution session for replans and +invocation revisions. Existing transforms run after nominal validation. A +replay-only source must return a non-replannable plan; the hook does not itself +supply Skill-specific semantic assembly or certify source capabilities. + Planning never steps simulation and never treats command completion as proof of a physical effect. @@ -69,6 +76,7 @@ contact or physical goal satisfaction. | Invocation, binding, and policies | `invocation.py`, `bindings.py`, `policies.py`, `control.py` | | Robot/task/scene state | `state.py`, `scene.py` | | Plans and runtime commands | `plans.py`, `runtime_commands.py` | +| Source-neutral trajectory inputs, schema and Skill semantic assembly | `trajectory_assembly.py`, `core.py` | | Session state machine | `execution.py` | | Verification request/result values | `verification.py` | | Runner and transports | `runner.py`, `transports.py`, `tracking.py` | @@ -203,3 +211,15 @@ 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. + +`AtomicActionEngine.build_plan_from_trajectory` admits full-joint motion through +an explicit versioned `SkillDescriptor.trajectory_contract`, then delegates to +the framework-owned `AtomicAction.assemble_trajectory_plan`. PickUp and +MoveJoints share their `_trajectory_semantics` hook between ordinary planning +and external motion; other Skills opt in explicitly. Skill code owns effects +and geometric facts, while the public builder owns commands and tracking. +`ActionPlan.trajectory_input` retains an owned input for same-grid rebuilding; +rebuilding reassembles semantics and rejects stale geometric-event facts. +Semantic stages/events do not grant expansion permissions. Input source/model +review remains provider-owned, and fixed replay remains non-replannable. +See [trajectory assembly details](execution.md#trajectory-backed-plan-assembly). diff --git a/agent_context/topics/atomic-actions/execution.md b/agent_context/topics/atomic-actions/execution.md index b2d4a34d4..0f9ab61e5 100644 --- a/agent_context/topics/atomic-actions/execution.md +++ b/agent_context/topics/atomic-actions/execution.md @@ -163,3 +163,25 @@ it never substitutes observed velocities or depends on peer active masks. It writes position and velocity targets through the existing `set_qpos` and `set_qvel` APIs. Inactive rows hold observed positions with zero target velocity. A nonzero damping gain is needed for the velocity error to contribute to tracking. + +## Trajectory-backed plan assembly + +`trajectory_assembly.py` owns immutable Skill schemas, timed inputs and semantic +assembly results. A concrete Skill opts in through `trajectory_contract` and +implements `_trajectory_semantics`; it does not override the public assembler. +The schema participates in the existing semantic catalog fingerprint. Registration +rejects a declared contract without an assembler, and unsupported Skills fail +before ordinary planning rather than falling back. + +The engine checks imported motion against restored initial state, position and +sampled/interval velocity limits, and bound endpoint claims. Sources retain +calibration, model identity, joint mapping and collision/physical qualification. +Stages are half-open; landmarks index actual samples. Shared stage layouts can +batch with ordered `env_ids`, but source admission may still require B=1. + +Same-grid rebuilding uses the plan's owned assembly input and current provenance +to regenerate commands, tracking, effects and monitor bounds. Supplied geometry +cannot survive a changed event pose without an explicit new input. Existing +non-joint plans and Skills without a contract retain their previous rebuild +path. Changed-grid execution, extra source recovery capabilities and phase +expansion permissions are separate contracts. diff --git a/agent_context/topics/data-pipeline/data-pipeline.md b/agent_context/topics/data-pipeline/data-pipeline.md index fe9cdf4e8..10d011bfa 100644 --- a/agent_context/topics/data-pipeline/data-pipeline.md +++ b/agent_context/topics/data-pipeline/data-pipeline.md @@ -14,6 +14,7 @@ belongs to [data-assets](../data-assets/data-assets.md). | Commit/discard/finalization dispatch | `embodichain/lab/gym/envs/managers/dataset_manager.py` | | Synchronous LeRobot persistence | `embodichain/lab/gym/envs/managers/datasets.py`: `LeRobotRecorder` | | Background persistence | `embodichain/lab/gym/envs/managers/async_datasets.py`: `AsyncLeRobotRecorder` | +| Real joint recordings and auditable normalization | `embodichain/data_pipeline/real2sim/`; optional external-file example in `examples/real2sim/` | | Depth sidecars | `embodichain/data_pipeline/depth_video/writer.py`, `reader.py` | | Offline collection retries and final partial batch | `embodichain/lab/scripts/run_env.py` | | File-only validation and grouped quality splits | `embodichain/data_pipeline/datasets/inspection.py`; `embodichain dataset validate\|split` | @@ -26,6 +27,20 @@ for fragment transactions, async ownership, depth output and shutdown. Read [offline dataset tools](offline-datasets.md) for lineage, quality filtering, validation, and conversion across Python/LeRobot versions. +## Real trajectory inputs + +`real2sim/contracts.py` owns dataset-independent, immutable source/joint/phase +values. `normalization.py` maps explicit units and joint names onto a supplied +full-robot layout; missing independent joints require declared holds, and mimic +values derive from drivers. The host supplies the control clock. Pure anchored +resampling lives in `compute/trajectory/anchored.py`; reports reference array +keys/hashes while the unit retains the actual bidirectional timing maps. +Unverified inputs may be normalized diagnostically but cannot enter +`motion.expansion.Real2SimSourceAdapter`. That adapter exports narrower free +permission windows to protect derivative stencils; full semantic phase ranges +remain in the normalized unit. Raw EEF evidence has explicit measured/derived +provenance. These contracts own no scene reconstruction, reset or persistence. + ## Online worker contract Shared image buffers are allocated by `init_rollout_buffer_from_config()` in @@ -127,3 +142,13 @@ Measure writer throughput using a runnable task, recording hardware, image size, environment count and writer settings alongside the results. The CPU-only `benchmark_segment_annotations.py` isolates the validated frame-map lookup from simulation and storage costs. + +`real2sim/bundle.py` persists raw source declarations and NPZ sidecars with +required content hashes, then reruns normalization at load. Formal admission +requires reviewed calibration and annotations. `real2sim/gripper.py` decodes an +explicit scalar/width convention through a versioned lookup; model-only estimates +remain unverified. Task Program source selection and Skill plan assembly belong +to `task_program/integrations`, not this data layer. +`gripper_release_duration` measures trailing release evidence from host-supplied +measured-open/contact masks; calibration and physical task acceptance stay with +the host, independently of source-input eligibility. diff --git a/agent_context/topics/motion-planning/motion-planning.md b/agent_context/topics/motion-planning/motion-planning.md index fb6c867dc..8ec22c2be 100644 --- a/agent_context/topics/motion-planning/motion-planning.md +++ b/agent_context/topics/motion-planning/motion-planning.md @@ -11,6 +11,8 @@ | Cartesian geometry and scalar timing | `embodichain/lab/sim/motion/planners/se3.py`, `bezier.py`, `_scalar_time_law.py`, `_blend_constraints.py` in that directory | | Collision planning and scene conversion | `embodichain/lab/sim/motion/planners/curobo/` | | NMG policy rollout and export metadata | `embodichain/lab/sim/motion/planners/neural_planner.py` | +| Anchored source-time resampling and audit maps | `embodichain/compute/trajectory/anchored.py` | +| Real2sim qpos source adaptation | `embodichain/lab/sim/motion/expansion/real2sim.py`; raw contracts belong to [data pipeline](../data-pipeline/data-pipeline.md) | | Pure interpolation, resampling and retiming | `embodichain/compute/trajectory/` | | Standalone playback and B=1 expansion host lifecycle | `embodichain/lab/sim/motion/execution.py` | | Expansion Profile loading, source adapters, candidate coordination and coverage | `embodichain/lab/sim/motion/expansion/` | @@ -62,7 +64,14 @@ do not establish continuous collision freedom or derivative compliance. ## Compute and execution ownership `compute.trajectory` owns pure interpolation, differentiation, resampling and -warping. New consumers import it directly; `lab.sim.utility.action_utils` keeps +warping. `anchored_batch.py` adds ragged tensor batches with explicit source and +anchor counts. CPU float64 mapping preparation shares the single-path algorithm; +`AnchoredBatchPlan.apply` reuses device maps through Torch or explicit Warp +CPU/CUDA kernels. Anchors stay exact, padding holds the last valid state, and +valid counts/masks exclude it from downstream evaluation. Allocation limits +cover the padded grid and output elements. The NumPy single-path API and the +B=1 registered source integration remain unchanged. New consumers import compute +directly; `lab.sim.utility.action_utils` keeps solver-dependent pose/IK adaptation and compatibility re-exports. Trapezoidal/ Double-S Warp profiles live in `compute/kinematics/_warp/trapezoidal.py`; compute must not import simulation objects. diff --git a/agent_context/topics/task-programs/task-programs.md b/agent_context/topics/task-programs/task-programs.md index c6950dffa..dfc0663a9 100644 --- a/agent_context/topics/task-programs/task-programs.md +++ b/agent_context/topics/task-programs/task-programs.md @@ -32,6 +32,7 @@ Semantic Calls and their scene/profile/effect contracts live inside | Compiled program model and AST expansion | `embodichain/lab/task_program/compiler/program.py` | | Semantic Call analysis and lowering | `embodichain/lab/task_program/compiler/lowering.py` | | Semantic Call execution | `embodichain/lab/task_program/runtime/executor.py`, `results.py` | +| Trusted nominal plan source hook | `runtime/expansion.py`: `TaskProgramNominalPlanFactory`; `integrations/environment.py`: runtime/bridge assembly | | Grounded call plan-transform request | `embodichain/lab/task_program/runtime/expansion.py` | | Parallel scheduling and safety boundary | `embodichain/lab/task_program/runtime/parallel.py`, `parallel_executor.py` | | Semantic Calls, scene, profile, effects, evidence | `embodichain/lab/task_program/semantics/` | @@ -80,6 +81,29 @@ Atomic plan transform; the bridge exposes immutable `TaskProgramExpansionRecord` provenance. Nonmatching calls and episodes without a profile keep the ordinary path. +A trusted `TaskProgramNominalPlanFactory` can provide a complete call-scoped +plan before ordinary Atomic planning. The engine snapshots and validates that +plan, then applies any existing plan transform. Execution sessions retain the +source during recovery and invocation revisions; source errors do not fall back +to another planner. `integrations/expert_sources.py` owns immutable versioned +factories, exact expanded call selections and capability preflight. The simulation +catalog fingerprints those declarations; configured task integrations select +`runtime_services.expert_source_providers` by stable IDs and content-pinned source +bundles. Bundle paths resolve from the integration file. Runtime assembly creates +fresh providers sharing the existing planning scene and embodiment profile. +Identity checks use weak references, so completed calls do not retain their +recordings; provider implementations must support weak references. Fixed-qpos +admission checks interval speeds as well as differentiated velocity references. + +The first registered source is B=1 fixed-qpos Pick replay. It uses the +[source-independent Skill trajectory assembler](../atomic-actions/execution.md#trajectory-backed-plan-assembly), +with the Skill owning semantic stages and pending effects. Skill trajectory +schemas are included in registration fingerprints. +Unverified input bundles, automatic replanning/retry policies, parallel segments +and configured combined expansion fail closed. Physical qualification and +receipt-backed source attachments remain separate host work. Do not insert +executable provider syntax into program YAML. + The configured combined transform supports batches and Pick Affordance expansion; timing expansion is not connected to this transform. Batch expansion records retain the stable environment IDs from the planning context, including diff --git a/docs/architecture/curated.json b/docs/architecture/curated.json index cf50a031a..bfcad6579 100644 --- a/docs/architecture/curated.json +++ b/docs/architecture/curated.json @@ -3499,9 +3499,9 @@ { "path": "embodichain/lab/task_program/integrations/environment.py", "symbol": "TaskProgramEnvironmentAdapter._assemble_execution_runtime", - "start_line": 768, - "end_line": 776, - "excerpt": " runtime = SemanticCallExecutor(\n semantic.compiler,\n observation_provider,\n command_sink,\n evidence_collector,\n clock=clock,\n runner_cfg=deepcopy(selected_runner_cfg),\n plan_transform_factory=plan_transform_factory,\n )" + "start_line": 804, + "end_line": 813, + "excerpt": " runtime = SemanticCallExecutor(\n semantic.compiler,\n observation_provider,\n command_sink,\n evidence_collector,\n clock=clock,\n runner_cfg=deepcopy(selected_runner_cfg),\n plan_transform_factory=plan_transform_factory,\n nominal_plan_factory=nominal_plan_factory,\n )" } ] }, @@ -3578,9 +3578,9 @@ { "path": "embodichain/lab/task_program/runtime/executor.py", "symbol": "SemanticCallExecutor._prepare_grounded_call", - "start_line": 969, - "end_line": 974, - "excerpt": " session = self._engine.start(\n (invocation,),\n context,\n eligible_mask=active_mask,\n plan_transform=plan_transform,\n )" + "start_line": 1108, + "end_line": 1114, + "excerpt": " session = self._engine.start(\n (invocation,),\n context,\n eligible_mask=active_mask,\n plan_transform=plan_transform,\n nominal_plan_provider=nominal_plan_provider,\n )" } ] }, diff --git a/docs/source/api_reference/atomic_trajectory_assembly.rst b/docs/source/api_reference/atomic_trajectory_assembly.rst new file mode 100644 index 000000000..c5451dde4 --- /dev/null +++ b/docs/source/api_reference/atomic_trajectory_assembly.rst @@ -0,0 +1,85 @@ +:orphan: + +Trajectory-backed Atomic Skill plans +=================================== + +Ordinary planning, admitted real recordings and trajectory candidates can share +a Skill-owned semantic assembler. The source supplies motion and annotations; +the Skill supplies expected effects and scene-monitor boundaries. The existing +plan builder lowers commands and tracking, and the existing execution runtime +retains effect assurance, recovery and state commits. + +Contracts and owned inputs +------------------------- + +.. currentmodule:: embodichain.lab.sim.atomic_actions + +.. autosummary:: + + SkillTrajectoryGeometry + SkillTrajectoryContract + SkillTrajectoryInput + SkillPlanSemantics + +.. automodule:: embodichain.lab.sim.atomic_actions.trajectory_assembly + :members: + +An installed ``SkillDescriptor.trajectory_contract`` explicitly declares the +accepted schema and revision, semantic stages, event IDs and optional geometric +facts type. Unsupported Skills reject trajectory assembly before ordinary +planning. The declaration is included in the Task Program registration +fingerprint. Geometry snapshots belong to the declared concrete type; they do +not establish measured physical evidence. + +Stages cover the execution path with half-open ranges. Event landmarks index +actual samples, so a stage-end event refers to ``stop - 1``. These semantic +annotations do not authorize expansion: existing ``TrajectoryPhase`` windows, +operators and joint permissions remain authoritative. + +Assembly and rebuilding +----------------------- + +.. automethod:: embodichain.lab.sim.atomic_actions.AtomicAction.assemble_trajectory_plan + :no-index: + +.. automethod:: embodichain.lab.sim.atomic_actions.AtomicActionEngine.build_plan_from_trajectory + :no-index: + +``build_plan_from_trajectory`` checks the restored start state, full-joint limits, +sampled and interval speeds, and endpoint authorization before Skill assembly. +This is discrete trajectory admission, not collision or physical-success proof. +Sources retain model/joint mapping, calibration, input review and host-specific +validation responsibilities. + +``ActionPlan.trajectory_input`` retains an owned assembly snapshot. The existing +``rebuild_plan_from_trajectory`` delegates supported plans to the same assembler, +rebuilding commands, tracking, semantic effects and dependency bounds together. +It still requires identical frame counts and arrival intervals. A candidate +that changes a geometric event while retaining old geometric facts is rejected; +the caller must create a new geometry family and assembly input. + +The original ``build_plan_from_recorded_trajectory`` engine method remains a +compatibility adapter using declared stage-end landmarks. It contains no +Pick-specific stage or effect logic. A replay-only source still requires a +fresh reset for another attempt and cannot enable recovery by changing an input +flag. + +Implemented Skills +------------------ + +PickUp and MoveJoints use the shared assembler for both ordinary planning and +external joint motion. PickUp derives an omitted expected grasp pose by FK at +the explicit grasp landmark, then declares an uncommitted held-object effect. +MoveJoints checks its terminal goal and ordered waypoint visits. Existing +non-joint command plans and Skills without an assembly contract retain their +original planning path. + +.. autoclass:: embodichain.lab.sim.atomic_actions.primitives.pick_up.PickTrajectoryGeometry + :members: + :no-index: + +The generic assembler supports batches sharing one stage layout and preserves +ordered ``env_ids``. The registered real2sim source remains B=1 fixed-qpos Pick; +the new Skill interface does not expand its admitted capabilities. Ragged +execution, changed-grid rebuilding, additional Skill assemblers and registered +source combined expansion require separate integration. diff --git a/docs/source/api_reference/embodichain/embodichain.compute.rst b/docs/source/api_reference/embodichain/embodichain.compute.rst index 3d098f798..3beb906a5 100644 --- a/docs/source/api_reference/embodichain/embodichain.compute.rst +++ b/docs/source/api_reference/embodichain/embodichain.compute.rst @@ -46,6 +46,9 @@ means trajectory deformation. differentiate_positions resample_in_time + AnchoredBatchPlan + prepare_anchored_batch + resample_with_anchors_batched interpolate_with_distance interpolate_with_nums resample_with_distance @@ -63,6 +66,14 @@ or motion limits. ``resample_in_time`` preserves endpoints and total duration while sampling the original time profile; callers must recompute derivatives after changing samples or timing. +Anchored batch resampling accepts ragged ``(B,N_max,D)`` paths with explicit +source lengths and anchor counts. ``prepare_anchored_batch`` computes auditable +float64 clocks on the CPU and prepares reusable device maps. Applying a plan +uses Torch by default or explicit Warp CPU/CUDA kernels, preserves float32 or +float64 qpos, and copies event samples exactly. Shorter outputs contain terminal +holds; use the plan's valid counts/mask to exclude padding. Mapping preparation +and interpolation do not check collisions or change the simulator clock. + Implementation modules ---------------------- @@ -78,6 +89,9 @@ Implementation modules .. automodule:: embodichain.compute.trajectory.warping :members: +.. automodule:: embodichain.compute.trajectory.anchored_batch + :members: + Kinematics API -------------- diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index c0ee80d99..e767ea556 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -716,6 +716,8 @@ embodichain.lab.sim.atomic_actions.control embodichain.lab.sim.atomic_actions.core --------------------------------------- +Trajectory-backed Skill assembly is described in :doc:`atomic_trajectory_assembly`. + .. currentmodule:: embodichain.lab.sim.atomic_actions.core .. autosummary:: diff --git a/docs/source/api_reference/real2sim_providers.rst b/docs/source/api_reference/real2sim_providers.rst new file mode 100644 index 000000000..21d0c6caf --- /dev/null +++ b/docs/source/api_reference/real2sim_providers.rst @@ -0,0 +1,47 @@ +:orphan: + +Registered real2sim source providers +=================================== + +These APIs admit content-pinned, non-executable source bundles and bind a nominal +source to one expanded Task Program call. Sources provide data and complete plans; +the existing host retains command execution, effect checks and persistence. +Review states describe input prerequisites and never imply physical success. + +Source bundles and gripper conventions +------------------------------------- + +.. automodule:: embodichain.data_pipeline.real2sim.bundle + :members: write_real2sim_bundle, load_real2sim_bundle + +.. automodule:: embodichain.data_pipeline.real2sim.gripper + :members: + +.. autoclass:: embodichain.data_pipeline.real2sim.GripperPositionCalibration + :no-index: + +.. autofunction:: embodichain.data_pipeline.real2sim.gripper_release_duration + :no-index: + +``gripper_release_duration`` measures trailing release evidence from host-supplied +measured-open and bilateral-contact masks using elapsed timestamps. The host +retains calibration, object containment/stability and task acceptance decisions. + +Registered nominal plans +----------------------- + +.. automodule:: embodichain.lab.task_program.integrations.expert_sources + :members: ExpertSourceProvider, ExpertSourceProviderFactory, ExpertSourceSelection, ExpertSourceRegistry + +.. automodule:: embodichain.lab.task_program.integrations.real2sim + :members: Real2SimExpertSourceProviderFactory, Real2SimExpertSourceProvider + +The initial registered source supports B=1, sequential, fixed-qpos Pick replay. +The provider calls the source-independent :doc:`atomic_trajectory_assembly` +interface; PickUp owns semantic segment and effect assembly. Replay-only sources reject +replanning/retry policies and configured combined expansion. Unverified source +bundles remain usable for offline diagnostics but cannot enter formal admission. +Provider instances must support weak references. Runtime identity tracking does +not retain completed providers or their recording arrays; active callbacks retain +their own provider. Both sampled velocity references and per-command-interval +speeds must satisfy the live robot limits, including at rapid reversals. diff --git a/docs/source/api_reference/real2sim_sources.rst b/docs/source/api_reference/real2sim_sources.rst new file mode 100644 index 000000000..9610c2fa5 --- /dev/null +++ b/docs/source/api_reference/real2sim_sources.rst @@ -0,0 +1,55 @@ +:orphan: + +Real2sim source foundations +========================== + +These APIs normalize explicit real joint recordings and adapt reviewed units to +the existing candidate pipeline. They do not reconstruct scenes or certify +physical task success. EEF evidence is distinguished from independently measured +poses; the initial normalizer implements qpos authority only. + +Data contracts +-------------- + +.. autosummary:: + + embodichain.data_pipeline.real2sim + +.. automodule:: embodichain.data_pipeline.real2sim + :members: EefPoseSamples, MimicJoint, Real2SimTrajectory, RobotJointLayout, SourcePhase + +.. automodule:: embodichain.data_pipeline.real2sim.contracts + :members: SourcePhase, MimicJoint, RobotJointLayout, Real2SimTrajectory, EefPoseSamples + +Normalization and timing +------------------------ + +.. automodule:: embodichain.data_pipeline.real2sim.normalization + :members: Real2SimNormalizationCfg, NormalizedReal2SimUnit, normalize_real2sim + +.. automodule:: embodichain.compute.trajectory.anchored + :members: AnchoredTimeMap, resample_with_anchors + +Candidate adaptation +-------------------- + +.. autoclass:: embodichain.lab.sim.motion.expansion.Real2SimSourceAdapter + :members: + +.. automodule:: embodichain.lab.sim.motion.expansion.real2sim + :members: Real2SimSourceAdapter + +Nominal plan integration +------------------------ + +A trusted nominal source supplies a complete ActionPlan before ordinary motion +planning. The engine still validates and snapshots its result. The session +retains the source during recovery and invocation revision. Existing plan +transforms run after nominal-plan validation. Registered source catalogs, +configured source services and real2sim-specific Pick assembly are described in +:doc:`real2sim_providers`; the hook itself does not certify their capabilities. + +.. autodata:: embodichain.lab.sim.atomic_actions.NominalPlanProvider + +.. automodule:: embodichain.lab.task_program.runtime.expansion + :members: TaskProgramNominalPlanFactory diff --git a/embodichain/compute/trajectory/__init__.py b/embodichain/compute/trajectory/__init__.py index 180765cc4..0dcd1edd1 100644 --- a/embodichain/compute/trajectory/__init__.py +++ b/embodichain/compute/trajectory/__init__.py @@ -21,6 +21,12 @@ from __future__ import annotations +from .anchored import AnchoredTimeMap, resample_with_anchors +from .anchored_batch import ( + AnchoredBatchPlan, + prepare_anchored_batch, + resample_with_anchors_batched, +) from .interpolation import interpolate_with_distance, interpolate_with_nums from .resampling import resample_with_distance from .timing import ( @@ -31,6 +37,11 @@ from .warping import sort_and_padding_key_frame, warp_trajectory_qpos __all__ = [ + "AnchoredTimeMap", + "AnchoredBatchPlan", + "prepare_anchored_batch", + "resample_with_anchors_batched", + "resample_with_anchors", "differentiate_positions", "resample_in_time", "retime_to_control_grid", diff --git a/embodichain/compute/trajectory/_warp/anchored.py b/embodichain/compute/trajectory/_warp/anchored.py new file mode 100644 index 000000000..c97687e48 --- /dev/null +++ b/embodichain/compute/trajectory/_warp/anchored.py @@ -0,0 +1,95 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Private float32/float64 anchored interpolation kernels.""" + +from __future__ import annotations + +from contextlib import nullcontext + +import torch +import warp as wp + +__all__ = [] + +# Keep float64 timing arithmetic comparable to the NumPy/Torch reference. +wp.set_module_options({"enable_backward": False, "fuse_fp": False}) + + +@wp.kernel +def _interpolate_f32( + q: wp.array3d(dtype=wp.float32), + left: wp.array2d(dtype=wp.int64), + right: wp.array2d(dtype=wp.int64), + weights: wp.array2d(dtype=wp.float64), + out: wp.array3d(dtype=wp.float32), +): + b, t, d = wp.tid() + lo = left[b, t] + hi = right[b, t] + if lo == hi: + out[b, t, d] = q[b, lo, d] + else: + difference = q[b, hi, d] - q[b, lo, d] + out[b, t, d] = wp.float32( + wp.float64(q[b, lo, d]) + weights[b, t] * wp.float64(difference) + ) + + +@wp.kernel +def _interpolate_f64( + q: wp.array3d(dtype=wp.float64), + left: wp.array2d(dtype=wp.int64), + right: wp.array2d(dtype=wp.int64), + weights: wp.array2d(dtype=wp.float64), + out: wp.array3d(dtype=wp.float64), +): + b, t, d = wp.tid() + lo = left[b, t] + hi = right[b, t] + if lo == hi: + out[b, t, d] = q[b, lo, d] + else: + out[b, t, d] = q[b, lo, d] + weights[b, t] * (q[b, hi, d] - q[b, lo, d]) + + +def evaluate( + positions: torch.Tensor, + left: torch.Tensor, + right: torch.Tensor, + weights: torch.Tensor, +) -> torch.Tensor: + """Launch on the caller's Torch stream; explicit backend errors propagate.""" + wp.init() + q = positions.contiguous() + out = positions.new_empty((*left.shape, positions.shape[-1])) + scope = ( + wp.ScopedStream(wp.stream_from_torch(q.device)) if q.is_cuda else nullcontext() + ) + with scope: + wp.launch( + kernel=_interpolate_f32 if q.dtype == torch.float32 else _interpolate_f64, + dim=out.shape, + inputs=[ + wp.from_torch(q), + wp.from_torch(left), + wp.from_torch(right), + wp.from_torch(weights), + ], + outputs=[wp.from_torch(out)], + device=str(q.device), + ) + return out diff --git a/embodichain/compute/trajectory/anchored.py b/embodichain/compute/trajectory/anchored.py new file mode 100644 index 000000000..a0e1f3ff6 --- /dev/null +++ b/embodichain/compute/trajectory/anchored.py @@ -0,0 +1,269 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Auditable fixed-clock resampling with exact source event anchors.""" + +from __future__ import annotations + +from dataclasses import dataclass +from collections.abc import Sequence +import math + +import numpy as np + +__all__ = ["AnchoredTimeMap", "resample_with_anchors"] + + +def _owned(array: np.ndarray) -> np.ndarray: + return np.frombuffer(array.tobytes(), dtype=array.dtype).reshape(array.shape) + + +@dataclass(frozen=True) +class AnchoredTimeMap: + """Bidirectional timing evidence and source interpolation brackets. + + ``mapped_source_times`` maps each original sample onto execution time; + ``source_query_times`` maps each execution sample back onto original time. + All times are relative seconds. Anchors refer to samples, not commands. + + Args: + source_times: Original zero-based clock. + mapped_source_times: Execution time assigned to each source sample. + execution_times: Zero-based destination control clock. + source_query_times: Original time queried by each destination sample. + left_indices: Left source interpolation bracket per destination sample. + right_indices: Right source interpolation bracket per destination sample. + weights: Interpolation fractions within the brackets. + anchor_source_indices: Ordered source samples that must be preserved. + anchor_execution_indices: Corresponding exact destination sample indices. + """ + + source_times: np.ndarray + mapped_source_times: np.ndarray + execution_times: np.ndarray + source_query_times: np.ndarray + left_indices: np.ndarray + right_indices: np.ndarray + weights: np.ndarray + anchor_source_indices: np.ndarray + anchor_execution_indices: np.ndarray + + def __post_init__(self) -> None: + for name in self.__dataclass_fields__: + value = np.asarray(getattr(self, name)) + if value.ndim != 1 or not np.isfinite(value).all(): + raise ValueError(f"{name} must be a finite one-dimensional array") + object.__setattr__(self, name, _owned(value)) + n, m = len(self.source_times), len(self.execution_times) + if n < 2 or m < 2 or len(self.mapped_source_times) != n: + raise ValueError("Time map needs nonempty source and execution clocks") + if any( + len(getattr(self, k)) != m + for k in ("source_query_times", "left_indices", "right_indices", "weights") + ): + raise ValueError("Execution mapping arrays must have matching lengths") + if ( + np.any(np.diff(self.source_times) <= 0) + or np.any(np.diff(self.execution_times) <= 0) + or np.any(np.diff(self.mapped_source_times) <= 0) + ): + raise ValueError("Source and mapped clocks must strictly increase") + for indices, bound in [ + (self.left_indices, n), + (self.right_indices, n), + (self.anchor_source_indices, n), + (self.anchor_execution_indices, m), + ]: + if ( + indices.dtype.kind not in "iu" + or np.any(indices < 0) + or np.any(indices >= bound) + ): + raise ValueError("Time map indices must be in-range integers") + if ( + len(self.anchor_source_indices) != len(self.anchor_execution_indices) + or np.any(self.left_indices > self.right_indices) + or np.any((self.weights < 0) | (self.weights > 1)) + ): + raise ValueError("Invalid interpolation or anchor mapping") + a, b = self.anchor_source_indices, self.anchor_execution_indices + if ( + len(a) < 2 + or a[0] != 0 + or a[-1] != n - 1 + or b[0] != 0 + or b[-1] != m - 1 + or np.any(a[1:] <= a[:-1]) + or np.any(b[1:] <= b[:-1]) + ): + raise ValueError( + "Anchors must be strictly ordered and include both endpoints" + ) + if ( + self.source_times[0] != 0 + or self.execution_times[0] != 0 + or np.any(np.diff(self.source_query_times) <= 0) + ): + raise ValueError( + "Relative source and execution clocks must start at zero and advance" + ) + reconstructed = self.source_times[self.left_indices] + self.weights * ( + self.source_times[self.right_indices] - self.source_times[self.left_indices] + ) + if ( + not np.allclose( + reconstructed, self.source_query_times, atol=1e-12, rtol=1e-10 + ) + or not np.allclose( + self.source_query_times[b], self.source_times[a], atol=1e-12, rtol=1e-10 + ) + or not np.allclose( + self.mapped_source_times[a], + self.execution_times[b], + atol=1e-12, + rtol=1e-10, + ) + ): + raise ValueError( + "Time mappings disagree with interpolation brackets or hard anchors" + ) + + +def resample_with_anchors( + positions: np.ndarray, + source_times: np.ndarray, + control_dt: float, + anchor_indices: Sequence[int], + *, + max_anchor_delay: float | None = None, + max_samples: int = 1_000_000, +) -> tuple[np.ndarray, AnchoredTimeMap]: + """Resample a path while preserving ordered hard event samples exactly. + + Anchors are rounded forward to control ticks. Between anchors a monotone + piecewise-affine time map retains source path timing. The first/last sample + are always anchors. If two events cannot fit within the permitted timing + error, fail instead of deleting a phase or silently extending its delay. + + Args: + positions: Finite floating ``(N,D)`` source positions. + source_times: Strictly increasing ``(N,)`` seconds, starting at zero. + control_dt: Positive destination control period; physics is untouched. + anchor_indices: Explicit hard source sample indices. + max_anchor_delay: Maximum delay per anchor; defaults to one control tick. + max_samples: Maximum output allocation, checked before creating the grid. + + Returns: + Resampled positions and an owned, bidirectional time/index mapping. + """ + q = np.asarray(positions) + t = np.asarray(source_times, dtype=np.float64) + if ( + q.ndim != 2 + or q.dtype.kind != "f" + or min(q.shape) < 1 + or len(q) < 2 + or not np.isfinite(q).all() + or t.shape != (len(q),) + or not np.isfinite(t).all() + or t[0] != 0 + or np.any(np.diff(t) <= 0) + ): + raise ValueError( + "Expected finite (N>=2,D) positions and a strictly increasing zero-based clock" + ) + mapping = _prepare_time_map( + t, + control_dt, + anchor_indices, + max_anchor_delay=max_anchor_delay, + max_samples=max_samples, + ) + result = q[mapping.left_indices] + mapping.weights[:, None] * ( + q[mapping.right_indices] - q[mapping.left_indices] + ) + result = result.astype(q.dtype) + result[mapping.anchor_execution_indices] = q[mapping.anchor_source_indices] + return result, mapping + + +def _prepare_time_map( + t: np.ndarray, + control_dt: float, + anchor_indices: Sequence[int], + *, + max_anchor_delay: float | None, + max_samples: int, +) -> AnchoredTimeMap: + """Build the common float64 clock without evaluating joint positions.""" + if ( + t.ndim != 1 + or len(t) < 2 + or not np.isfinite(t).all() + or t[0] != 0 + or np.any(np.diff(t) <= 0) + ): + raise ValueError("Expected a strictly increasing zero-based source clock") + if ( + type(control_dt) not in (float, int) + or not math.isfinite(control_dt) + or control_dt <= 0 + ): + raise ValueError("control_dt must be positive and finite") + if type(max_samples) is not int or max_samples < 2: + raise ValueError("max_samples must be an integer >= 2") + delay = float(control_dt) if max_anchor_delay is None else max_anchor_delay + if type(delay) not in (float, int) or not math.isfinite(delay) or delay < 0: + raise ValueError("max_anchor_delay must be finite and nonnegative") + raw_anchors = tuple(anchor_indices) + if any(type(i) is not int or not 0 <= i < len(t) for i in raw_anchors): + raise ValueError("Anchors must be in-range integer sample indices") + anchors = np.array(sorted({0, len(t) - 1, *raw_anchors}), dtype=np.int64) + ticks = [] + for source_index in anchors: + ratio = float(t[source_index] / control_dt) + if not math.isfinite(ratio) or ratio >= max_samples: + raise ValueError("Execution grid exceeds max_samples") + nearest = round(ratio) + tick = ( + nearest + if math.isclose(ratio, nearest, rel_tol=0.0, abs_tol=1e-8) + else math.ceil(ratio) + ) + tick = max(ticks[-1] + 1, tick) if ticks else 0 + if tick >= max_samples: + raise ValueError("Execution grid exceeds max_samples") + if tick * control_dt - t[source_index] > delay + 1e-10: + raise ValueError( + "Required anchors cannot fit the control clock within max_anchor_delay" + ) + ticks.append(tick) + target_anchors = np.array(ticks, dtype=np.int64) + anchor_times = target_anchors * control_dt + grid = np.arange(ticks[-1] + 1, dtype=np.float64) * control_dt + mapped = np.interp(t, t[anchors], anchor_times) + query = np.interp(grid, anchor_times, t[anchors]) + upper = np.searchsorted(t, query, side="right").clip(1, len(t) - 1) + lower = upper - 1 + weights = ((query - t[lower]) / (t[upper] - t[lower])).clip(0.0, 1.0) + # Exact source samples are represented as copies, independent of backend. + query[target_anchors] = t[anchors] + lower[target_anchors] = anchors + upper[target_anchors] = anchors + weights[target_anchors] = 0.0 + return AnchoredTimeMap( + t, mapped, grid, query, lower, upper, weights, anchors, target_anchors + ) diff --git a/embodichain/compute/trajectory/anchored_batch.py b/embodichain/compute/trajectory/anchored_batch.py new file mode 100644 index 000000000..dbbe29a1c --- /dev/null +++ b/embodichain/compute/trajectory/anchored_batch.py @@ -0,0 +1,318 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Prepare auditable ragged clocks once and evaluate batched joint paths.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Literal + +import numpy as np +import torch + +from .anchored import AnchoredTimeMap, _prepare_time_map + +__all__ = [ + "AnchoredBatchPlan", + "prepare_anchored_batch", + "resample_with_anchors_batched", +] + + +@dataclass(frozen=True) +class AnchoredBatchPlan: + """Reusable interpolation maps on one execution device. + + Host rows retain immutable float64 timing evidence. Device buffers contain + only interpolation indices/weights; no robot or trajectory values are owned. + Padding copies the last valid source sample and is excluded by valid_counts. + Plans can be shared by trajectories with the same clocks and source lengths. + + Args: + rows: Nonempty tuple of independently validated source/execution maps. + device: Device used by every tensor passed to :meth:`apply`. + max_samples: Bound on the total rectangular grid size B * M_max. + """ + + rows: tuple[AnchoredTimeMap, ...] + device: torch.device | str = "cpu" + max_samples: int = 1_000_000 + _left: torch.Tensor = field(init=False, repr=False, compare=False) + _right: torch.Tensor = field(init=False, repr=False, compare=False) + _weights: torch.Tensor = field(init=False, repr=False, compare=False) + _source_lengths: torch.Tensor = field(init=False, repr=False, compare=False) + _valid_counts: torch.Tensor = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + rows = tuple(self.rows) + if not rows or any(not isinstance(r, AnchoredTimeMap) for r in rows): + raise ValueError("rows must contain nonempty AnchoredTimeMap values") + if type(self.max_samples) is not int or self.max_samples < 2: + raise ValueError("max_samples must be an integer >= 2") + width = max(len(row.execution_times) for row in rows) + if len(rows) * width > self.max_samples: + raise ValueError("Padded execution grid exceeds max_samples") + device = torch.device(self.device) + if device.type not in ("cpu", "cuda"): + raise ValueError("Anchored batches support CPU and CUDA devices") + left = np.empty((len(rows), width), dtype=np.int64) + right = np.empty_like(left) + weights = np.zeros((len(rows), width), dtype=np.float64) + for b, row in enumerate(rows): + size = len(row.execution_times) + left[b] = right[b] = len(row.source_times) - 1 + left[b, :size] = row.left_indices + right[b, :size] = row.right_indices + weights[b, :size] = row.weights + # A valid external map may encode a knot as an interval with weight + # one. Canonicalize hard anchors to copies for every backend. + events = row.anchor_execution_indices + left[b, events] = right[b, events] = row.anchor_source_indices + weights[b, events] = 0.0 + object.__setattr__(self, "rows", rows) + for name, value in ( + ("_left", left), + ("_right", right), + ("_weights", weights), + ("_source_lengths", np.array([len(r.source_times) for r in rows])), + ("_valid_counts", np.array([len(r.execution_times) for r in rows])), + ): + object.__setattr__(self, name, torch.tensor(value, device=device)) + object.__setattr__(self, "device", self._left.device) + + @property + def source_lengths(self) -> torch.Tensor: + """Return an owned integer (B,) tensor of valid source sample counts.""" + return self._source_lengths.clone() + + @property + def valid_counts(self) -> torch.Tensor: + """Return an owned integer (B,) tensor of valid execution sample counts.""" + return self._valid_counts.clone() + + @property + def valid_mask(self) -> torch.Tensor: + """Return a boolean (B,M_max) mask excluding terminal-hold padding.""" + return ( + torch.arange(self._left.shape[1], device=self.device)[None] + < self._valid_counts[:, None] + ) + + def apply( + self, + positions: torch.Tensor, + *, + backend: Literal["torch", "warp"] = "torch", + max_output_elements: int = 100_000_000, + ) -> torch.Tensor: + """Evaluate a batch using prepared clocks, preserving dtype and device. + + Args: + positions: Float32/float64 (B,N_max,D) tensor. Only each row's + declared valid source prefix must be finite; padding is ignored. + backend: Explicit execution backend. Warp errors propagate without + silently selecting another backend. Torch supports qpos autograd; + Warp rejects gradient-bearing inputs in this implementation. + max_output_elements: Bound on B * M_max * D before output allocation. + + Returns: + Owned (B,M_max,D) tensor with exact anchor copies and terminal holds. + + Raises: + ValueError: Invalid shape, device, dtype, backend or finite samples. + """ + if backend not in ("torch", "warp"): + raise ValueError("backend must be torch or warp") + if ( + not isinstance(positions, torch.Tensor) + or positions.ndim != 3 + or positions.dtype not in (torch.float32, torch.float64) + or positions.shape[0] != len(self.rows) + or positions.shape[2] < 1 + or positions.device != self.device + or positions.shape[1] < max(len(r.source_times) for r in self.rows) + ): + raise ValueError( + "Expected float32/float64 (B,N_max,D) positions on the plan device" + ) + if type(max_output_elements) is not int or max_output_elements < 1: + raise ValueError("max_output_elements must be a positive integer") + if self._left.numel() * positions.shape[-1] > max_output_elements: + raise ValueError("Resampled positions exceed max_output_elements") + source_mask = ( + torch.arange(positions.shape[1], device=self.device)[None] + < self._source_lengths[:, None] + ) + if not (torch.isfinite(positions).all(dim=-1) | ~source_mask).all(): + raise ValueError("Valid source positions must be finite") + if backend == "warp": + if positions.requires_grad: + raise ValueError( + "Warp anchored resampling does not support autograd; use torch" + ) + from ._warp.anchored import evaluate + + return evaluate(positions, self._left, self._right, self._weights) + dimensions = positions.shape[-1] + q0 = positions.gather(1, self._left[..., None].expand(-1, -1, dimensions)) + q1 = positions.gather(1, self._right[..., None].expand(-1, -1, dimensions)) + # Match the single NumPy reference: subtract in qpos dtype, accumulate + # with float64 time weights, then cast once. Never demote float64 qpos. + interpolated = (q0.double() + self._weights[..., None] * (q1 - q0).double()).to( + positions.dtype + ) + return torch.where((self._left == self._right)[..., None], q0, interpolated) + + +def prepare_anchored_batch( + source_times: torch.Tensor, + control_dt: float, + anchor_indices: torch.Tensor, + *, + source_lengths: torch.Tensor, + anchor_counts: torch.Tensor, + device: torch.device | str | None = None, + max_anchor_delay: float | None = None, + max_samples: int = 1_000_000, +) -> AnchoredBatchPlan: + """Prepare independent ragged clocks and reusable device interpolation maps. + + Mapping/anchor decisions run on the CPU in float64, using exactly the same + implementation as the single-path API. Valid prefixes alone are inspected. + Preparation may synchronize CUDA metadata once; apply reuses device buffers. + + Args: + source_times: Floating (B,N_max) relative seconds; each valid prefix + starts at zero and strictly increases. Padding may contain NaN. + control_dt: Shared positive control period in seconds; physics is untouched. + anchor_indices: Integer (B,K_max) source indices. Unused entries are ignored. + source_lengths: Integer (B,) valid source counts, each >= 2. + anchor_counts: Integer (B,) valid anchor counts, possibly zero. First and + last source samples are always added, and duplicate anchors coalesce. + device: Evaluation device; defaults to source_times.device. + max_anchor_delay: Permitted forward event delay; None means one control tick. + max_samples: Bound on both per-row and total padded execution samples. + + Returns: + A plan retaining per-row audit maps, device buffers and valid lengths. + + Raises: + ValueError: Invalid inputs or any row whose events cannot fit the clock. + Row-level errors identify the batch row; no failed row is discarded. + """ + if ( + not isinstance(source_times, torch.Tensor) + or source_times.ndim != 2 + or source_times.dtype not in (torch.float32, torch.float64) + or source_times.shape[0] < 1 + ): + raise ValueError("source_times must be floating (B,N_max), B >= 1") + batch, width = source_times.shape + if type(max_samples) is not int or max_samples < 2 * batch: + raise ValueError("max_samples must allow at least two samples per batch row") + for name, value, shape in ( + ("source_lengths", source_lengths, (batch,)), + ("anchor_counts", anchor_counts, (batch,)), + ): + if ( + not isinstance(value, torch.Tensor) + or value.shape != shape + or value.dtype not in (torch.int32, torch.int64) + ): + raise ValueError(f"{name} must be an integer (B,) tensor") + if ( + not isinstance(anchor_indices, torch.Tensor) + or anchor_indices.ndim != 2 + or anchor_indices.shape[0] != batch + or anchor_indices.dtype not in (torch.int32, torch.int64) + ): + raise ValueError("anchor_indices must be an integer (B,K_max) tensor") + lengths = source_lengths.detach().cpu().tolist() + counts = anchor_counts.detach().cpu().tolist() + if any(n < 2 or n > width for n in lengths): + raise ValueError("source_lengths must be between 2 and N_max") + if any(k < 0 or k > anchor_indices.shape[1] for k in counts): + raise ValueError("anchor_counts must be between 0 and K_max") + times = source_times.detach().to(device="cpu", dtype=torch.float64).numpy() + anchors = anchor_indices.detach().cpu().numpy() + rows = [] + for b, (n, k) in enumerate(zip(lengths, counts)): + try: + rows.append( + _prepare_time_map( + times[b, :n], + control_dt, + anchors[b, :k].tolist(), + max_anchor_delay=max_anchor_delay, + max_samples=max_samples // batch, + ) + ) + except ValueError as error: + raise ValueError(f"Batch row {b}: {error}") from error + return AnchoredBatchPlan( + tuple(rows), source_times.device if device is None else device, max_samples + ) + + +def resample_with_anchors_batched( + positions: torch.Tensor, + source_times: torch.Tensor, + control_dt: float, + anchor_indices: torch.Tensor, + *, + source_lengths: torch.Tensor, + anchor_counts: torch.Tensor, + backend: Literal["torch", "warp"] = "torch", + max_anchor_delay: float | None = None, + max_samples: int = 1_000_000, + max_output_elements: int = 100_000_000, +) -> tuple[torch.Tensor, AnchoredBatchPlan]: + """Resample ragged joint paths onto fixed clocks with exact event samples. + + Args: + positions: Float32/float64 (B,N_max,D) joint paths. + source_times: Floating (B,N_max) relative seconds; padding is ignored. + control_dt: Shared destination control period, in seconds. + anchor_indices: Integer (B,K_max) hard source sample indices. + source_lengths: Integer (B,) valid source counts. + anchor_counts: Integer (B,) valid anchor counts. + backend: Explicit torch (default) or warp execution backend. + max_anchor_delay: Maximum forward anchor delay; defaults to one tick. + max_samples: Maximum padded execution grid size B * M_max. + max_output_elements: Maximum output allocation B * M_max * D. + + Returns: + Resampled positions and reusable plan. Use plan.valid_counts or + plan.valid_mask to exclude padding. This does not compute qvel or + qualify physical motion; those remain caller responsibilities. + """ + if not isinstance(positions, torch.Tensor): + raise ValueError("positions must be a tensor") + plan = prepare_anchored_batch( + source_times, + control_dt, + anchor_indices, + source_lengths=source_lengths, + anchor_counts=anchor_counts, + device=positions.device, + max_anchor_delay=max_anchor_delay, + max_samples=max_samples, + ) + return ( + plan.apply(positions, backend=backend, max_output_elements=max_output_elements), + plan, + ) diff --git a/embodichain/data_pipeline/__init__.py b/embodichain/data_pipeline/__init__.py index 744419060..f6fc7a07b 100644 --- a/embodichain/data_pipeline/__init__.py +++ b/embodichain/data_pipeline/__init__.py @@ -27,9 +27,9 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from . import datasets, depth_video, engine, recording + from . import datasets, depth_video, engine, recording, real2sim -__all__ = ["datasets", "depth_video", "engine", "recording"] +__all__ = ["datasets", "depth_video", "engine", "recording", "real2sim"] def __getattr__(name: str) -> ModuleType: diff --git a/embodichain/data_pipeline/real2sim/__init__.py b/embodichain/data_pipeline/real2sim/__init__.py new file mode 100644 index 000000000..e7f8cc07f --- /dev/null +++ b/embodichain/data_pipeline/real2sim/__init__.py @@ -0,0 +1,39 @@ +# ---------------------------------------------------------------------------- +# 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 real-trajectory data contracts, independent of dataset formats.""" + +from __future__ import annotations + +from .contracts import ( + EefPoseSamples, + MimicJoint, + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, +) + +from .gripper import GripperPositionCalibration, gripper_release_duration + +__all__ = [ + "GripperPositionCalibration", + "gripper_release_duration", + "EefPoseSamples", + "MimicJoint", + "Real2SimTrajectory", + "RobotJointLayout", + "SourcePhase", +] diff --git a/embodichain/data_pipeline/real2sim/bundle.py b/embodichain/data_pipeline/real2sim/bundle.py new file mode 100644 index 000000000..047996e0d --- /dev/null +++ b/embodichain/data_pipeline/real2sim/bundle.py @@ -0,0 +1,222 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Content-checked, non-executable real2sim source bundles.""" + +from __future__ import annotations + +from dataclasses import asdict, fields +import hashlib +import io +import json +import zipfile +from pathlib import Path + +import numpy as np + +from .contracts import ( + EefPoseSamples, + MimicJoint, + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, + _plain, +) +from .normalization import ( + Real2SimNormalizationCfg, + NormalizedReal2SimUnit, + normalize_real2sim, +) + +__all__ = ["write_real2sim_bundle", "load_real2sim_bundle"] + + +def _digest(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _verified_bytes(path: Path, expected: str, limit: int, name: str) -> bytes: + # Decode the same immutable bytes that were hashed, even if a caller replaces + # an external file concurrently with admission. + with path.open("rb") as stream: + data = stream.read(limit + 1) + if len(data) > limit or hashlib.sha256(data).hexdigest() != expected: + raise ValueError(f"Source {name} digest mismatch or oversized content") + return data + + +def _closed(value: object, keys: set[str], name: str) -> dict: + if not isinstance(value, dict) or set(value) != keys: + raise ValueError(f"{name} must contain exactly {sorted(keys)}") + return value + + +def write_real2sim_bundle( + source: Real2SimTrajectory, + layout: RobotJointLayout, + cfg: Real2SimNormalizationCfg, + directory: str | Path, +) -> Path: + """Write raw evidence and explicit normalization settings into a new directory. + + Args: + source: Owned raw source values; review states are preserved verbatim. + layout: Complete target joint model and restored initial state. + cfg: Explicit normalization settings. + directory: New output directory, separate from the original recording. + + Returns: + Path to the bundle JSON; its SHA-256 pins the entire declaration. + """ + normalize_real2sim(source, layout, cfg) + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=False) + arrays = { + "positions": source.positions, + "timestamps": source.timestamps, + "initial_positions": layout.initial_positions, + } + if source.velocities is not None: + arrays["velocities"] = source.velocities + data = { + f.name: _plain(getattr(source, f.name)) + for f in fields(source) + if f.name not in {"positions", "timestamps", "velocities", "eef", "phases"} + } + data["phases"] = [asdict(p) for p in source.phases] + data["has_velocities"] = source.velocities is not None + data["eef"] = None + if source.eef is not None: + arrays["eef_poses"] = source.eef.poses + data["eef"] = { + f.name: getattr(source.eef, f.name) + for f in fields(source.eef) + if f.name != "poses" + } + for name in ("reference_to_base", "link_to_tcp"): + data["eef"][name] = data["eef"][name].tolist() + target = { + f.name: getattr(layout, f.name) + for f in fields(layout) + if f.name not in {"initial_positions", "mimics"} + } + target["mimics"] = [asdict(m) for m in layout.mimics] + options = {f.name: getattr(cfg, f.name) for f in fields(cfg)} + np.savez_compressed(directory / "source.npz", **arrays) + payload = { + "schema_version": 1, + "arrays": "source.npz", + "arrays_sha256": _digest(directory / "source.npz"), + "source": data, + "layout": target, + "normalization": options, + } + path = directory / "bundle.json" + path.write_text( + json.dumps(payload, sort_keys=True, allow_nan=False, indent=2) + "\n" + ) + return path + + +def load_real2sim_bundle( + path: str | Path, *, expected_sha256: str, require_verified: bool = True +) -> NormalizedReal2SimUnit: + """Validate bundle content and rerun normalization before any simulation. + + Args: + path: Bundle JSON file with a local NPZ sidecar. + expected_sha256: Required content digest from trusted integration config. + require_verified: Require reviewed prerequisites for formal execution. + False allows offline inspection only. + + Returns: + Fresh owned normalized values; no simulator or executable imports are loaded. + """ + path = Path(path) + declaration = _verified_bytes(path, expected_sha256, 1_000_000, "bundle") + payload = _closed( + json.loads(declaration), + { + "schema_version", + "arrays", + "arrays_sha256", + "source", + "layout", + "normalization", + }, + "bundle", + ) + if payload["schema_version"] != 1 or payload["arrays"] != "source.npz": + raise ValueError("Unsupported bundle version or sidecar path") + sidecar = path.parent / payload["arrays"] + array_bytes = _verified_bytes( + sidecar, payload["arrays_sha256"], 256_000_000, "array" + ) + raw = _closed( + payload["source"], + {f.name for f in fields(Real2SimTrajectory)} + - {"positions", "timestamps", "velocities"} + | {"has_velocities"}, + "source", + ) + layout = _closed( + payload["layout"], + {f.name for f in fields(RobotJointLayout)} - {"initial_positions"}, + "layout", + ) + options = _closed( + payload["normalization"], + {f.name for f in fields(Real2SimNormalizationCfg)}, + "normalization", + ) + with zipfile.ZipFile(io.BytesIO(array_bytes)) as archive: + if sum(info.file_size for info in archive.infolist()) > 256_000_000: + raise ValueError("Source arrays exceed the uncompressed size limit") + if type(raw["has_velocities"]) is not bool: + raise ValueError("has_velocities must be boolean") + with np.load(io.BytesIO(array_bytes), allow_pickle=False) as arrays: + expected = {"positions", "timestamps", "initial_positions"} + if raw["has_velocities"]: + expected.add("velocities") + if raw["eef"] is not None: + expected.add("eef_poses") + if set(arrays.files) != expected: + raise ValueError("Source bundle array keys do not match declarations") + eef = raw["eef"] + if eef is not None: + _closed(eef, {f.name for f in fields(EefPoseSamples)} - {"poses"}, "EEF") + eef = EefPoseSamples(poses=arrays["eef_poses"], **eef) + source = Real2SimTrajectory( + **{ + k: v + for k, v in raw.items() + if k not in {"phases", "has_velocities", "eef"} + }, + positions=arrays["positions"], + timestamps=arrays["timestamps"], + velocities=arrays["velocities"] if raw["has_velocities"] else None, + phases=tuple(SourcePhase(**p) for p in raw["phases"]), + eef=eef, + ) + target = RobotJointLayout( + **{k: v for k, v in layout.items() if k != "mimics"}, + initial_positions=arrays["initial_positions"], + mimics=tuple(MimicJoint(**m) for m in layout["mimics"]), + ) + unit = normalize_real2sim(source, target, Real2SimNormalizationCfg(**options)) + if require_verified and not unit.eligible: + raise ValueError("Source bundle prerequisites remain unverified") + return unit diff --git a/embodichain/data_pipeline/real2sim/contracts.py b/embodichain/data_pipeline/real2sim/contracts.py new file mode 100644 index 000000000..ef5eab7ea --- /dev/null +++ b/embodichain/data_pipeline/real2sim/contracts.py @@ -0,0 +1,442 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Dataset-independent real-trajectory values; no simulator or dataset imports.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from collections.abc import Mapping +from types import MappingProxyType +import hashlib +import json +from typing import Literal + +import numpy as np + +__all__ = [ + "SourcePhase", + "MimicJoint", + "RobotJointLayout", + "Real2SimTrajectory", + "EefPoseSamples", +] + + +def _text(value: str, name: str) -> str: + if type(value) is not str or not value.strip() or value != value.strip(): + raise ValueError(f"{name} must be a nonempty stripped string") + return value + + +def _names(values: tuple[str, ...], name: str) -> tuple[str, ...]: + names = tuple(values) + for item in names: + _text(item, name) + if len(set(names)) != len(names): + raise ValueError(f"{name} must contain unique names") + return names + + +def _array(value: np.ndarray, shape: tuple[int | None, ...], name: str) -> np.ndarray: + array = np.asarray(value) + if ( + array.ndim != len(shape) + or any(n is not None and array.shape[i] != n for i, n in enumerate(shape)) + or array.dtype.kind not in "fiu" + or not np.isfinite(array).all() + ): + raise ValueError(f"{name} must be a finite real array with shape {shape}") + # Bytes-backed storage cannot be made writable again by a caller. + return np.frombuffer(array.tobytes(order="C"), dtype=array.dtype).reshape( + array.shape + ) + + +def _json(value: Mapping[str, object]) -> Mapping[str, object]: + def freeze(item): + if isinstance(item, dict): + return MappingProxyType({k: freeze(v) for k, v in item.items()}) + if isinstance(item, list): + return tuple(freeze(v) for v in item) + return item + + return freeze(json.loads(json.dumps(_plain(value), allow_nan=False))) + + +def _plain(value): + if isinstance(value, Mapping): + return {k: _plain(v) for k, v in value.items()} + if isinstance(value, tuple): + return [_plain(v) for v in value] + return value + + +def _transform(value: np.ndarray, name: str) -> np.ndarray: + matrix = _array(value, (4, 4), name) + if ( + not np.allclose(matrix[3], [0, 0, 0, 1], atol=1e-8) + or not np.allclose(matrix[:3, :3].T @ matrix[:3, :3], np.eye(3), atol=1e-6) + or not np.isclose(np.linalg.det(matrix[:3, :3]), 1.0, atol=1e-6) + ): + raise ValueError(f"{name} must be a rigid transform") + return matrix + + +@dataclass(frozen=True) +class SourcePhase: + """An explicit half-open source-sample interval and augmentation permission. + + Args: + name: Unique phase identifier within this source unit. + start: Inclusive first sample index. + stop: Exclusive final sample index. + kind: Free motion, contact or hold annotation. + allowed_operators: Explicit permissions, available only in free phases. + """ + + name: str + start: int + stop: int + kind: Literal["free", "contact", "hold"] + allowed_operators: tuple[str, ...] = () + + def __post_init__(self) -> None: + _text(self.name, "phase name") + if ( + type(self.start) is not int + or type(self.stop) is not int + or not 0 <= self.start < self.stop + ): + raise ValueError("Phase requires integer 0 <= start < stop") + if self.kind not in ("free", "contact", "hold"): + raise ValueError("Unknown phase kind") + operators = _names(self.allowed_operators, "allowed operators") + if self.kind != "free" and operators: + raise ValueError("Contact and hold phases cannot authorize operators") + object.__setattr__(self, "allowed_operators", operators) + + +@dataclass(frozen=True) +class MimicJoint: + """Canonical-unit relation ``joint = multiplier * parent + offset``. + + Args: + joint: Dependent target joint name. + parent: Driver target joint name. + multiplier: Signed coupling ratio. + offset: Offset in the target joint's canonical unit. + """ + + joint: str + parent: str + multiplier: float = 1.0 + offset: float = 0.0 + + def __post_init__(self) -> None: + _text(self.joint, "mimic joint") + _text(self.parent, "mimic parent") + if ( + self.joint == self.parent + or not np.isfinite([self.multiplier, self.offset]).all() + ): + raise ValueError("Invalid mimic relation") + + +@dataclass(frozen=True) +class RobotJointLayout: + """Full target joint order and explicit initial-state/augmentation ownership. + + Args: + robot_id: Target robot identity, matching the recorded embodiment. + robot_revision: Exact model/calibration revision required for qpos replay. + joint_names: Full simulator joint order. + joint_units: Canonical rad/m units in that order. + initial_positions: Restored initial state used for compatibility checks. + expandable_joint_names: Independent observed joints authorized for expansion. + mimics: Acyclic dependent-joint relations in canonical units. + """ + + robot_id: str + robot_revision: str + joint_names: tuple[str, ...] + joint_units: tuple[str, ...] + initial_positions: np.ndarray + expandable_joint_names: tuple[str, ...] = () + mimics: tuple[MimicJoint, ...] = () + + def __post_init__(self) -> None: + _text(self.robot_id, "robot_id") + _text(self.robot_revision, "robot_revision") + names = _names(self.joint_names, "joint_names") + units = tuple(self.joint_units) + if ( + not names + or len(units) != len(names) + or any(u not in ("rad", "m") for u in units) + ): + raise ValueError("Target layout requires one rad/m unit per named joint") + expandable = _names(self.expandable_joint_names, "expandable_joint_names") + mimics = tuple(self.mimics) + if not all(isinstance(m, MimicJoint) for m in mimics): + raise TypeError("mimics must contain MimicJoint values") + mimic_names = _names(tuple(m.joint for m in mimics), "mimic joints") + if not set(expandable) <= set(names) - set(mimic_names): + raise ValueError("Only independent target joints may be expandable") + available = set(names) - set(mimic_names) + remaining = list(mimics) + while remaining: + ready = [m for m in remaining if m.parent in available and m.joint in names] + if not ready: + raise ValueError("Mimic graph has an unknown joint/parent or a cycle") + for m in ready: + if units[names.index(m.joint)] != units[names.index(m.parent)]: + raise ValueError("Mimic joints must share their parent's unit") + available.add(m.joint) + remaining.remove(m) + object.__setattr__(self, "joint_names", names) + object.__setattr__(self, "joint_units", units) + object.__setattr__(self, "expandable_joint_names", expandable) + object.__setattr__(self, "mimics", mimics) + object.__setattr__( + self, + "initial_positions", + _array(self.initial_positions, (len(names),), "initial_positions"), + ) + + +@dataclass(frozen=True) +class EefPoseSamples: + """Optional EEF evidence aligned to the source joint timestamps in v1. + + ``reference_to_base`` maps reference-frame points into the robot base; + ``link_to_tcp`` is the TCP pose in the observed link frame. Pose positions + use ``position_unit``; transform translations always use metres. Independent + EEF clock alignment must be performed explicitly before creating this value. + + Args: + poses: Source-aligned ``(N,7)`` positions and quaternions. + reference_frame: Frame in which source poses are expressed. + observed_link: Link whose pose was recorded. + tcp_id: Target tool frame after applying the declared link transform. + reference_to_base: Rigid reference-to-robot-base transform. + link_to_tcp: Rigid TCP pose expressed in the observed link frame. + origin: Measurement lineage, kept separate from numerical consistency. + quaternion_order: Explicit quaternion component convention. + position_unit: Source pose translation unit. + """ + + poses: np.ndarray + reference_frame: str + observed_link: str + tcp_id: str + reference_to_base: np.ndarray + link_to_tcp: np.ndarray + origin: Literal["measured", "fk_derived", "transformed"] + quaternion_order: Literal["xyzw", "wxyz"] = "xyzw" + position_unit: Literal["m", "mm"] = "m" + + def __post_init__(self) -> None: + for name in ("reference_frame", "observed_link", "tcp_id"): + _text(getattr(self, name), name) + if self.origin not in ("measured", "fk_derived", "transformed"): + raise ValueError("Unknown EEF evidence origin") + if self.quaternion_order not in ("xyzw", "wxyz") or self.position_unit not in ( + "m", + "mm", + ): + raise ValueError("EEF units and quaternion order must be explicit") + object.__setattr__(self, "poses", _array(self.poses, (None, 7), "EEF poses")) + object.__setattr__( + self, + "reference_to_base", + _transform(self.reference_to_base, "reference_to_base"), + ) + object.__setattr__( + self, "link_to_tcp", _transform(self.link_to_tcp, "link_to_tcp") + ) + + +@dataclass(frozen=True) +class Real2SimTrajectory: + """Owned real-robot samples, explicit units and reviewed semantic annotations. + + Raw arrays retain their dtype; optional velocities use joint units per second. + Neither calibration nor annotation status + defaults to verified. This value owns no simulator, task runner or dataset. + + Args: + source_id: Dataset or source identity. + source_revision: Immutable source version or file digest. + unit_id: Selected recording or semantic slice identifier. + robot_id: Recorded embodiment identity. + robot_revision: Recorded robot model/calibration revision. + joint_names: Source column names, before target-order mapping. + joint_units: Source units for each column; decode gripper encodings first. + positions: Recorded ``(N,D)`` joint positions. + timestamps: Strictly increasing source clock values. + time_unit: Explicit source clock unit. + calibration_revision: Version of the calibration evidence. + annotation_revision: Version of the phase annotations. + calibration_status: Reviewed input calibration state. + annotation_status: Reviewed annotation state. + phases: Ordered, nonoverlapping semantic intervals; gaps remain locked. + velocities: Optional raw joint velocities in joint units per second. + eef: Optional source-aligned EEF evidence. + provenance: JSON-compatible preprocessing, correction and source metadata. + """ + + source_id: str + source_revision: str + unit_id: str + robot_id: str + robot_revision: str + joint_names: tuple[str, ...] + joint_units: tuple[str, ...] + positions: np.ndarray + timestamps: np.ndarray + time_unit: Literal["s", "ms", "us", "ns"] + calibration_revision: str + annotation_revision: str + calibration_status: Literal["verified", "unverified", "rejected"] = "unverified" + annotation_status: Literal["verified", "unverified", "rejected"] = "unverified" + phases: tuple[SourcePhase, ...] = () + velocities: np.ndarray | None = None + eef: EefPoseSamples | None = None + provenance: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + for name in ( + "source_id", + "source_revision", + "unit_id", + "robot_id", + "robot_revision", + "calibration_revision", + "annotation_revision", + ): + _text(getattr(self, name), name) + for status in (self.calibration_status, self.annotation_status): + if status not in ("verified", "unverified", "rejected"): + raise ValueError("Unknown verification status") + names = _names(self.joint_names, "joint_names") + units = tuple(self.joint_units) + if ( + not names + or len(units) != len(names) + or any(u not in ("rad", "deg", "m", "mm") for u in units) + ): + raise ValueError( + "Every source joint requires an explicit supported unit; decode grippers upstream" + ) + positions = _array(self.positions, (None, len(names)), "positions") + n = len(positions) + if n < 2: + raise ValueError("A source unit requires at least two samples") + timestamps = _array(self.timestamps, (n,), "timestamps") + if self.time_unit not in ("s", "ms", "us", "ns"): + raise ValueError("Unknown timestamp unit") + if not (timestamps[1:] > timestamps[:-1]).all(): + raise ValueError( + "Source timestamps must strictly increase; coalesce duplicates explicitly upstream" + ) + phases = tuple(self.phases) + if not all(isinstance(p, SourcePhase) for p in phases): + raise TypeError("phases must contain SourcePhase values") + _names(tuple(p.name for p in phases), "phase names") + previous = 0 + for phase in phases: + if phase.start < previous or phase.stop > n: + raise ValueError( + "Phases must be ordered, nonoverlapping and within the source" + ) + previous = phase.stop + if self.eef is not None and ( + not isinstance(self.eef, EefPoseSamples) or len(self.eef.poses) != n + ): + raise ValueError( + "v1 EEF evidence must align with all source joint timestamps" + ) + object.__setattr__(self, "joint_names", names) + object.__setattr__(self, "joint_units", units) + object.__setattr__(self, "positions", positions) + object.__setattr__(self, "timestamps", timestamps) + object.__setattr__(self, "phases", phases) + object.__setattr__( + self, + "velocities", + ( + None + if self.velocities is None + else _array(self.velocities, positions.shape, "velocities") + ), + ) + object.__setattr__(self, "provenance", _json(self.provenance)) + + @property + def seconds(self) -> np.ndarray: + """Elapsed source time in seconds, subtracting integer epochs first.""" + scale = {"s": 1.0, "ms": 0.001, "us": 1e-6, "ns": 1e-9}[self.time_unit] + return (self.timestamps - self.timestamps[0]).astype(np.float64) * scale + + def fingerprint(self) -> str: + """Return a content digest including data, units, annotations and calibration. + + Returns: + Stable SHA-256 hex digest over arrays and explicitly supplied provenance. + """ + digest = hashlib.sha256() + metadata = { + k: getattr(self, k) + for k in ( + "source_id", + "source_revision", + "unit_id", + "robot_id", + "robot_revision", + "joint_names", + "joint_units", + "time_unit", + "calibration_revision", + "annotation_revision", + "calibration_status", + "annotation_status", + ) + } + metadata["phases"] = [vars(p) for p in self.phases] + metadata["provenance"] = _plain(self.provenance) + arrays = {"positions": self.positions, "timestamps": self.timestamps} + if self.velocities is not None: + arrays["velocities"] = self.velocities + if self.eef is not None: + metadata["eef"] = { + k: getattr(self.eef, k) + for k in ( + "reference_frame", + "observed_link", + "tcp_id", + "origin", + "quaternion_order", + "position_unit", + ) + } + for k in ("poses", "reference_to_base", "link_to_tcp"): + arrays["eef_" + k] = getattr(self.eef, k) + digest.update(json.dumps(metadata, sort_keys=True, allow_nan=False).encode()) + for name, array in sorted(arrays.items()): + digest.update(json.dumps([name, array.dtype.str, array.shape]).encode()) + digest.update(array.tobytes(order="C")) + return digest.hexdigest() diff --git a/embodichain/data_pipeline/real2sim/gripper.py b/embodichain/data_pipeline/real2sim/gripper.py new file mode 100644 index 000000000..8954132d6 --- /dev/null +++ b/embodichain/data_pipeline/real2sim/gripper.py @@ -0,0 +1,149 @@ +# ---------------------------------------------------------------------------- +# 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 scalar-to-driver calibration, independent of a gripper model.""" + +from __future__ import annotations + +from dataclasses import dataclass +import numpy as np + +__all__ = ["GripperPositionCalibration", "gripper_release_duration"] + + +def gripper_release_duration( + timestamps: np.ndarray, + gripper_open: np.ndarray, + bilateral_contact: np.ndarray, +) -> float: + """Measure continuous trailing release evidence in a recorded trace. + + The host computes gripper_open from measured positions and its explicit + calibration/tolerance. Release requires an open gripper without bilateral + grasp contact. This measures elapsed time, not a count of samples, and does + not itself establish object stability, task success or source eligibility. + + Args: + timestamps: Finite, strictly increasing (N,) elapsed seconds. + gripper_open: Boolean (N,) measured-open evidence, not controller targets. + bilateral_contact: Boolean (N,) simultaneous left/right grasp contacts. + + Returns: + Seconds since the start of the final uninterrupted release interval. + Empty, singleton or currently unreleased traces return zero. + + Raises: + ValueError: Invalid shapes, masks, timestamps or nonfinite intervals. + """ + times = np.asarray(timestamps, dtype=np.float64) + opened = np.asarray(gripper_open) + contact = np.asarray(bilateral_contact) + if ( + times.ndim != 1 + or opened.shape != times.shape + or contact.shape != times.shape + or opened.dtype != np.bool_ + or contact.dtype != np.bool_ + or not np.isfinite(times).all() + ): + raise ValueError( + "Release evidence requires finite (N,) times and boolean masks" + ) + intervals = np.diff(times) + if not np.isfinite(intervals).all() or np.any(intervals <= 0): + raise ValueError( + "Release timestamps must strictly increase with finite intervals" + ) + released = opened & ~contact + if len(times) < 2 or not released[-1]: + return 0.0 + interruptions = np.flatnonzero(~released) + start = int(interruptions[-1] + 1) if len(interruptions) else 0 + return float(times[-1] - times[start]) + + +@dataclass(frozen=True) +class GripperPositionCalibration: + """Decode a declared width convention through a reviewed lookup table. + + Args: + encoding: One of normalized_closure, normalized_open or width_m. + max_width_m: Width corresponding to a fully open source gripper. + widths_m: Increasing measured or model-derived jaw widths in metres. + driver_positions_rad: Corresponding independent driver angles. + calibration_id: Versioned evidence identity. + evidence: Source/calibration reference; never inferred from sample range. + status: Review state; model-only estimates remain unverified. + """ + + encoding: str + max_width_m: float + widths_m: tuple[float, ...] + driver_positions_rad: tuple[float, ...] + calibration_id: str + evidence: str + status: str = "unverified" + + def __post_init__(self) -> None: + if self.encoding not in {"normalized_closure", "normalized_open", "width_m"}: + raise ValueError("Gripper encoding must be explicitly supported") + widths = tuple(float(x) for x in self.widths_m) + angles = tuple(float(x) for x in self.driver_positions_rad) + if ( + not np.isfinite(self.max_width_m) + or self.max_width_m <= 0 + or len(widths) < 2 + or len(widths) != len(angles) + or not np.isfinite([*widths, *angles]).all() + or np.any(np.diff(widths) <= 0) + or widths[0] > 0 + or widths[-1] < self.max_width_m + ): + raise ValueError("Calibration table must cover [0, max_width_m]") + if not (np.all(np.diff(angles) >= 0) or np.all(np.diff(angles) <= 0)): + raise ValueError("Gripper driver calibration must be monotone") + if any( + not isinstance(x, str) or not x.strip() + for x in (self.calibration_id, self.evidence) + ): + raise ValueError("Calibration ID and evidence are required") + if self.status not in {"verified", "unverified", "rejected"}: + raise ValueError("Unknown gripper calibration status") + object.__setattr__(self, "widths_m", widths) + object.__setattr__(self, "driver_positions_rad", angles) + + def decode(self, values: np.ndarray) -> np.ndarray: + """Return driver radians without clipping malformed source values. + + Args: + values: Finite scalar source samples under the declared encoding. + + Returns: + Driver angles with the same shape as the source samples. + """ + if self.status == "rejected": + raise ValueError("Rejected gripper calibration cannot decode samples") + values = np.asarray(values, dtype=np.float64) + upper = self.max_width_m if self.encoding == "width_m" else 1.0 + if not np.isfinite(values).all() or np.any((values < 0) | (values > upper)): + raise ValueError("Gripper samples exceed their declared encoding domain") + width = ( + values + if self.encoding == "width_m" + else self.max_width_m + * (1 - values if self.encoding == "normalized_closure" else values) + ) + return np.interp(width, self.widths_m, self.driver_positions_rad) diff --git a/embodichain/data_pipeline/real2sim/normalization.py b/embodichain/data_pipeline/real2sim/normalization.py new file mode 100644 index 000000000..f45460ac5 --- /dev/null +++ b/embodichain/data_pipeline/real2sim/normalization.py @@ -0,0 +1,475 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Auditable qpos normalization using explicit robot-layout and phase contracts.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import MISSING, dataclass +import hashlib +import json +import math +from typing import Literal + +import numpy as np +import torch + +from embodichain.utils import configclass +from embodichain.compute.trajectory import ( + AnchoredTimeMap, + resample_with_anchors, + differentiate_positions, +) +from .contracts import ( + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, + _array, + _json, + _plain, +) + +__all__ = ["Real2SimNormalizationCfg", "NormalizedReal2SimUnit", "normalize_real2sim"] + + +@configclass +class Real2SimNormalizationCfg: + """Explicit source-to-host normalization settings. + + Args: + control_dt: Host-owned command interval in seconds; required. + authority: Source authority. This first implementation supports qpos only. + joint_name_map: Source-to-target aliases. Unlisted names map identically. + hold_joint_names: Missing independent joints explicitly held at the supplied initial state. + initial_tolerance: Initial-position tolerance in each target joint's canonical unit. + mimic_tolerance: Redundant mimic observation tolerance in canonical joint units. + max_anchor_delay: Permitted forward anchor timing error; None means one control tick. + max_samples: Maximum normalized output length. + allow_unverified: Allow diagnostic normalization, never verified execution eligibility. + """ + + control_dt: float = MISSING + authority: Literal["qpos", "eef_pose", "hybrid"] = "qpos" + joint_name_map: dict[str, str] = {} + hold_joint_names: tuple[str, ...] = () + initial_tolerance: float = 1e-4 + mimic_tolerance: float = 1e-5 + max_anchor_delay: float | None = None + max_samples: int = 1_000_000 + allow_unverified: bool = False + + +@dataclass(frozen=True) +class NormalizedReal2SimUnit: + """Owned canonical trajectory plus original evidence and bidirectional timing. + + Args: + source: Original immutable source evidence. + joint_names: Canonical full robot joint order. + joint_units: Canonical rad/m units. + positions: Normalized ``(T,D)`` full-joint positions. + velocities: Velocities recomputed on the complete normalized path. + dt: Per-sample seconds, with zero at the first sample. + phases: Full semantic phases on the normalized grid. + expandable_joint_indices: Explicit independent-joint expansion permissions. + time_map: Source/execution time and index correspondence. + eef_poses: Optional base-frame TCP positions and xyzw quaternions. + report: Immutable normalization evidence; physical success is not implied. + """ + + source: Real2SimTrajectory + joint_names: tuple[str, ...] + joint_units: tuple[str, ...] + positions: np.ndarray + velocities: np.ndarray + dt: np.ndarray + phases: tuple[SourcePhase, ...] + expandable_joint_indices: tuple[int, ...] + time_map: AnchoredTimeMap + eef_poses: np.ndarray | None + report: Mapping[str, object] + + def __post_init__(self) -> None: + names = tuple(self.joint_names) + if not names or len(set(names)) != len(names): + raise ValueError("Normalized joint names must be nonempty and unique") + positions = _array(self.positions, (None, len(names)), "normalized positions") + n = len(positions) + if ( + n < 2 + or not isinstance(self.source, Real2SimTrajectory) + or not isinstance(self.time_map, AnchoredTimeMap) + ): + raise ValueError("Invalid normalized unit source or time map") + dt = _array(self.dt, (n,), "normalized dt") + if dt[0] != 0 or np.any(dt[1:] <= 0) or len(self.time_map.execution_times) != n: + raise ValueError( + "Normalized intervals require zero first dt and positive later dt" + ) + if not np.allclose( + np.cumsum(dt), self.time_map.execution_times, atol=1e-10, rtol=1e-10 + ): + raise ValueError("Normalized dt does not match the execution clock") + if len(self.joint_units) != len(names) or any( + u not in ("rad", "m") for u in self.joint_units + ): + raise ValueError("Normalized units must be rad/m for every joint") + if any( + type(i) is not int or not 0 <= i < len(names) + for i in self.expandable_joint_indices + ): + raise ValueError("Invalid expandable joint indices") + if len(set(self.expandable_joint_indices)) != len( + self.expandable_joint_indices + ): + raise ValueError("Expandable joint indices must be unique") + previous = 0 + for phase in self.phases: + if ( + not isinstance(phase, SourcePhase) + or phase.start < previous + or phase.stop > n + ): + raise ValueError("Invalid normalized phase partition") + previous = phase.stop + object.__setattr__(self, "joint_names", names) + object.__setattr__(self, "joint_units", tuple(self.joint_units)) + object.__setattr__(self, "positions", positions) + object.__setattr__( + self, + "velocities", + _array(self.velocities, positions.shape, "normalized velocities"), + ) + object.__setattr__(self, "dt", dt) + object.__setattr__(self, "phases", tuple(self.phases)) + object.__setattr__( + self, "expandable_joint_indices", tuple(self.expandable_joint_indices) + ) + object.__setattr__( + self, + "eef_poses", + ( + None + if self.eef_poses is None + else _array(self.eef_poses, (n, 7), "normalized EEF") + ), + ) + object.__setattr__(self, "report", _json(self.report)) + + @property + def revision(self) -> str: + """Content-derived normalization revision used by SourceContext.""" + return str(self.report["normalized_revision"]) + + @property + def eligible(self) -> bool: + """Whether reviewed input prerequisites passed, not physical qualification.""" + return self.report["input_prerequisites_verified"] is True + + def report_dict(self) -> dict: + """Return an independent JSON-compatible normalization report.""" + return _plain(self.report) + + +def _normalize_eef( + source: Real2SimTrajectory, mapping: AnchoredTimeMap +) -> np.ndarray | None: + eef = source.eef + if eef is None: + return None + from scipy.spatial.transform import Rotation, Slerp + + poses = eef.poses.astype(np.float64) + quat = poses[:, 3:] + if eef.quaternion_order == "wxyz": + quat = quat[:, [1, 2, 3, 0]] + norm = np.linalg.norm(quat, axis=1) + if not np.allclose(norm, 1.0, atol=1e-4, rtol=0): + raise ValueError("EEF quaternions must be normalized") + matrices = np.tile(np.eye(4), (len(poses), 1, 1)) + matrices[:, :3, :3] = Rotation.from_quat(quat).as_matrix() + matrices[:, :3, 3] = poses[:, :3] * (0.001 if eef.position_unit == "mm" else 1.0) + canonical = eef.reference_to_base @ matrices @ eef.link_to_tcp + q = mapping.source_query_times + positions = np.stack( + [np.interp(q, source.seconds, canonical[:, i, 3]) for i in range(3)], axis=1 + ) + rotations = Slerp(source.seconds, Rotation.from_matrix(canonical[:, :3, :3]))( + q + ).as_quat() + for i in range(1, len(rotations)): + if rotations[i - 1] @ rotations[i] < 0: + rotations[i] *= -1 + return np.c_[positions, rotations] + + +def normalize_real2sim( + source: Real2SimTrajectory, + layout: RobotJointLayout, + cfg: Real2SimNormalizationCfg, +) -> NormalizedReal2SimUnit: + """Normalize one qpos-authoritative unit without live robot or scene access. + + Args: + source: Immutable raw trajectory and evidence declarations. + layout: Full target robot joint order, initial state and mimic graph. + cfg: Explicit host clock, aliases and missing-joint policy. + + Returns: + Canonical full-joint positions/velocities, optional base-frame TCP evidence, + exact event anchors, source/execution mappings and an auditable report. + Eligibility only establishes reviewed input prerequisites. + """ + if not isinstance(source, Real2SimTrajectory) or not isinstance( + layout, RobotJointLayout + ): + raise TypeError("source and layout must use the real2sim value contracts") + if not isinstance(cfg, Real2SimNormalizationCfg): + raise TypeError("cfg must be Real2SimNormalizationCfg") + cfg.validate() + if cfg.authority not in ("qpos", "eef_pose", "hybrid"): + raise ValueError("Unknown real2sim authority") + if cfg.authority != "qpos": + raise NotImplementedError( + "eef_pose and hybrid require an explicit IK/replanning provider" + ) + if type(cfg.allow_unverified) is not bool: + raise ValueError("allow_unverified must be boolean") + for name in ("initial_tolerance", "mimic_tolerance"): + value = getattr(cfg, name) + if type(value) not in (float, int) or not math.isfinite(value) or value < 0: + raise ValueError(f"{name} must be finite and nonnegative") + if (source.robot_id, source.robot_revision) != ( + layout.robot_id, + layout.robot_revision, + ): + raise ValueError("qpos authority requires the same robot ID and revision") + statuses = (source.calibration_status, source.annotation_status) + if "rejected" in statuses: + raise ValueError("Rejected calibration or annotations cannot be normalized") + eligible = all(s == "verified" for s in statuses) and bool(source.phases) + if not eligible and not cfg.allow_unverified: + raise ValueError( + "Unverified calibration/phase annotations require explicit diagnostic mode" + ) + aliases = dict(cfg.joint_name_map) + if not set(aliases) <= set(source.joint_names): + raise ValueError("Joint aliases include unknown source names") + mapped = tuple(aliases.get(n, n) for n in source.joint_names) + if len(set(mapped)) != len(mapped) or not set(mapped) <= set(layout.joint_names): + raise ValueError("Joint aliases must map uniquely into the target layout") + hold = tuple(cfg.hold_joint_names) + mimic_names = {m.joint for m in layout.mimics} + if ( + len(set(hold)) != len(hold) + or not set(hold) <= set(layout.joint_names) - set(mapped) - mimic_names + ): + raise ValueError("Hold policy must name only missing independent joints") + missing = set(layout.joint_names) - set(mapped) - mimic_names - set(hold) + if missing: + raise ValueError( + f"Missing independent joints without explicit hold policy: {sorted(missing)}" + ) + if not set(layout.expandable_joint_names) <= set(mapped) - mimic_names: + raise ValueError("Only observed independent joints may be expanded") + n, d = len(source.positions), len(layout.joint_names) + full = np.empty((n, d), dtype=np.float64) + factors = { + "rad": ("rad", 1.0), + "deg": ("rad", np.pi / 180), + "m": ("m", 1.0), + "mm": ("m", 0.001), + } + column_mapping = [] + for column, target in enumerate(mapped): + target_index = layout.joint_names.index(target) + unit, scale = factors[source.joint_units[column]] + if unit != layout.joint_units[target_index]: + raise ValueError(f"Joint dimension mismatch for {target}") + full[:, target_index] = source.positions[:, column] * scale + column_mapping.append( + { + "source": source.joint_names[column], + "target": target, + "target_index": target_index, + "scale": scale, + } + ) + held = {} + for name in hold: + index = layout.joint_names.index(name) + full[:, index] = layout.initial_positions[index] + held[name] = float(layout.initial_positions[index]) + available = set(mapped) - mimic_names | set(hold) + remaining = list(layout.mimics) + resolved_mimics = [] + while remaining: + ready = [m for m in remaining if m.parent in available] + for mimic in ready: + j = layout.joint_names.index(mimic.joint) + parent = layout.joint_names.index(mimic.parent) + derived = full[:, parent] * mimic.multiplier + mimic.offset + if mimic.joint in mapped and not np.allclose( + full[:, j], derived, atol=cfg.mimic_tolerance, rtol=0 + ): + raise ValueError( + f"Observed mimic joint {mimic.joint} conflicts with its driver" + ) + full[:, j] = derived + available.add(mimic.joint) + remaining.remove(mimic) + resolved_mimics.append(mimic) + if not ready: + raise ValueError("Mimic graph could not resolve its independent drivers") + error = np.abs(full[0] - layout.initial_positions) + if np.any(error > cfg.initial_tolerance): + raise ValueError( + "Source initial positions do not match the declared restored initial state" + ) + anchors = {0, n - 1} + for phase in source.phases: + anchors.update((phase.start, phase.stop - 1)) + positions, time_map = resample_with_anchors( + full, + source.seconds, + cfg.control_dt, + sorted(anchors), + max_anchor_delay=cfg.max_anchor_delay, + max_samples=cfg.max_samples, + ) + for mimic in resolved_mimics: + positions[:, layout.joint_names.index(mimic.joint)] = ( + positions[:, layout.joint_names.index(mimic.parent)] * mimic.multiplier + + mimic.offset + ) + index_map = dict( + zip( + time_map.anchor_source_indices.tolist(), + time_map.anchor_execution_indices.tolist(), + ) + ) + phases = tuple( + SourcePhase( + p.name, + index_map[p.start], + index_map[p.stop - 1] + 1, + p.kind, + p.allowed_operators, + ) + for p in source.phases + ) + dt = np.r_[0.0, np.full(len(positions) - 1, cfg.control_dt)] + velocities = differentiate_positions( + torch.from_numpy(positions[None].copy()), torch.from_numpy(dt[None]) + ).numpy()[0] + eef = _normalize_eef(source, time_map) + report = { + "schema_version": 1, + "normalization_version": "qpos-anchors-v1", + "authority": "qpos", + "source_digest": source.fingerprint(), + "source_id": source.source_id, + "source_revision": source.source_revision, + "unit_id": source.unit_id, + "source_description": { + "joint_names": source.joint_names, + "joint_units": source.joint_units, + "time_unit": source.time_unit, + "velocity_time_unit": "s", + "provenance": _plain(source.provenance), + "eef": ( + None + if source.eef is None + else { + "reference_frame": source.eef.reference_frame, + "observed_link": source.eef.observed_link, + "tcp_id": source.eef.tcp_id, + "reference_to_base": source.eef.reference_to_base.tolist(), + "link_to_tcp": source.eef.link_to_tcp.tolist(), + "origin": source.eef.origin, + "quaternion_order": source.eef.quaternion_order, + "position_unit": source.eef.position_unit, + } + ), + }, + "robot_id": layout.robot_id, + "robot_revision": layout.robot_revision, + "joint_names": layout.joint_names, + "joint_units": layout.joint_units, + "column_mapping": column_mapping, + "held_positions": held, + "expandable_joint_names": layout.expandable_joint_names, + "normalization_options": { + "initial_tolerance": cfg.initial_tolerance, + "mimic_tolerance": cfg.mimic_tolerance, + "max_anchor_delay": cfg.max_anchor_delay, + "max_samples": cfg.max_samples, + }, + "mimics": [vars(m) for m in layout.mimics], + "initial_max_error": float(error.max()), + "control_dt": cfg.control_dt, + "source_samples": n, + "execution_samples": len(positions), + "source_duration_s": float(source.seconds[-1]), + "execution_duration_s": float(time_map.execution_times[-1]), + "max_anchor_delay_s": float( + np.max( + time_map.execution_times[time_map.anchor_execution_indices] + - source.seconds[time_map.anchor_source_indices] + ) + ), + "calibration_revision": source.calibration_revision, + "calibration_status": source.calibration_status, + "annotation_revision": source.annotation_revision, + "annotation_status": source.annotation_status, + "input_prerequisites_verified": eligible, + "eef_origin": None if source.eef is None else source.eef.origin, + "independent_eef_measurement": source.eef is not None + and source.eef.origin == "measured", + "source_phases": [vars(p) for p in source.phases], + "execution_phases": [vars(p) for p in phases], + "unannotated_interphase_samples": "locked", + "qvel_policy": "recomputed_after_full_path_resampling", + "time_mapping_arrays": { + name: { + "key": "map_" + name, + "shape": list(getattr(time_map, name).shape), + "dtype": getattr(time_map, name).dtype.str, + "sha256": hashlib.sha256(getattr(time_map, name).tobytes()).hexdigest(), + } + for name in time_map.__dataclass_fields__ + }, + } + digest = hashlib.sha256( + json.dumps(report, sort_keys=True, allow_nan=False).encode() + ) + for array in (positions, velocities, dt, *(() if eef is None else (eef,))): + digest.update(array.tobytes()) + report["normalized_revision"] = "real2sim-v1:" + digest.hexdigest() + return NormalizedReal2SimUnit( + source, + layout.joint_names, + layout.joint_units, + positions, + velocities, + dt, + phases, + tuple(layout.joint_names.index(n) for n in layout.expandable_joint_names), + time_map, + eef, + report, + ) diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index 45dea3326..2301db904 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -68,7 +68,7 @@ ) from .core import AtomicAction, ObjectSemantics, SkillDescriptor from .effects import StateDelta -from .engine import AtomicActionEngine, PlanTransform +from .engine import AtomicActionEngine, PlanTransform, NominalPlanProvider from .execution import ( ExecutionEvent, ExecutionEventKind, @@ -130,6 +130,12 @@ ) from .transports import EndpointCommandRouter, EndpointCommandTransport from .trajectory_adapter import ActionPlanTemplateAdapter +from .trajectory_assembly import ( + SkillPlanSemantics, + SkillTrajectoryContract, + SkillTrajectoryGeometry, + SkillTrajectoryInput, +) from .tracking import ( BASE_POSE_CHANNEL, JOINT_POSITION_CHANNEL, @@ -287,6 +293,7 @@ "AtomicAction", "AtomicActionEngine", "PlanTransform", + "NominalPlanProvider", "BUILTIN_ACTION_TYPES", "BATCH_INVERSE_KINEMATICS_CAPABILITY", "CARTESIAN_POSE_CAPABILITY", @@ -426,6 +433,10 @@ "SceneSnapshotSupplier", "SceneEntityPose", "SkillDescriptor", + "SkillPlanSemantics", + "SkillTrajectoryContract", + "SkillTrajectoryGeometry", + "SkillTrajectoryInput", "SkillBindingContract", "SkillEndpointRequirement", "SkillResourceSlot", diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index c2c8fa00c..cdad25d34 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -63,6 +63,11 @@ TrackingFrame, TrackingSetpoint, ) +from .trajectory_assembly import ( + SkillPlanSemantics, + SkillTrajectoryContract, + SkillTrajectoryInput, +) if TYPE_CHECKING: from embodichain.lab.sim.objects import Robot @@ -147,6 +152,8 @@ class SkillDescriptor: """Whether completion reports motion execution without physical-effect proof.""" binding_contract: SkillBindingContract | None = None """Explicit generic resource contract used by Task Program lowering.""" + trajectory_contract: SkillTrajectoryContract | None = None + """Explicit source-neutral trajectory assembly schema, when supported.""" def __post_init__(self) -> None: if not isinstance(self.skill_id, str) or not self.skill_id: @@ -170,6 +177,13 @@ def __post_init__(self) -> None: "SkillDescriptor.binding_contract must be a " "SkillBindingContract or None." ) + if ( + self.trajectory_contract is not None + and type(self.trajectory_contract) is not SkillTrajectoryContract + ): + raise TypeError( + "trajectory_contract must be SkillTrajectoryContract or None" + ) class AtomicAction(Generic[GoalT, OptionsT], ABC): @@ -203,6 +217,9 @@ class AtomicAction(Generic[GoalT, OptionsT], ABC): does not silently expose a new skill identifier. """ + trajectory_contract: ClassVar[SkillTrajectoryContract | None] = None + """Optional explicit opt-in to source-independent joint-plan assembly.""" + def __init_subclass__(cls, **kwargs: Any) -> None: """Reject skill classes that bypass framework-owned scene binding.""" super().__init_subclass__(**kwargs) @@ -211,6 +228,11 @@ def __init_subclass__(cls, **kwargs: Any) -> None: "AtomicAction subclasses must implement _plan(); the public " "plan() method is framework-owned." ) + if "assemble_trajectory_plan" in cls.__dict__: + raise TypeError( + "AtomicAction subclasses must implement _trajectory_semantics(); " + "assemble_trajectory_plan() is framework-owned." + ) def __init__( self, @@ -300,6 +322,7 @@ def descriptor(cls) -> SkillDescriptor: agent_visible=cls.agent_visible, open_loop=cls.open_loop, binding_contract=cls.__dict__.get("binding_contract"), + trajectory_contract=cls.__dict__.get("trajectory_contract"), ) def resolve_request( @@ -485,6 +508,96 @@ def _scene_dependencies( """Return scene entities whose poses materially affect this plan.""" return collect_scene_dependencies(request.goal) + def assemble_trajectory_plan( + self, + request: ResolvedActionRequest[GoalT, OptionsT], + context: PlanningContext, + motion: SkillTrajectoryInput, + ) -> ActionPlan: + """Build a complete plan using the Skill's source-independent semantics. + + Args: + request: Grounded Skill invocation with authoritative endpoint bindings. + context: Current observed robot, task and scene snapshot. + motion: Owned timed motion, stages and event landmarks. + + Returns: + Commands, tracking and uncommitted effects assembled together. + + Raises: + ValueError: The Skill has not opted in or an input violates its schema. + """ + self.require_goal(request) + contract = self.descriptor().trajectory_contract + if contract is None: + raise ValueError( + f"Skill {self.skill_id!r} does not support trajectory assembly" + ) + if type(motion) is not SkillTrajectoryInput: + raise TypeError("motion must be exactly SkillTrajectoryInput") + owned = motion.snapshot() + contract.validate_input(owned) + timed = owned.trajectory + if ( + timed.batch_size != context.batch_size + or timed.robot_dof != context.robot.robot_dof + ): + raise ValueError("Trajectory dimensions must match the planning context") + if timed.positions.device != context.robot.qpos.device: + raise ValueError("Trajectory and planning context must share a device") + if timed.positions.dtype != context.robot.qpos.dtype: + raise ValueError("Trajectory and planning context must share a joint dtype") + if not torch.equal(timed.env_ids, context.env_ids): + raise ValueError("Trajectory must preserve ordered planning env_ids") + ticks = timed.dt / context.require_control_dt() + if not torch.isclose(ticks, ticks.round(), atol=1e-5, rtol=0).all(): + raise ValueError( + "Trajectory arrival intervals must align to the control grid" + ) + names = {s.name for s in owned.segments} or {self.skill_id} + if any(gate.segment_name not in names for gate in request.phase_effect_gates): + raise ValueError("Trajectory is missing a required phase-effect gate stage") + semantics = self._trajectory_semantics(request, context, owned) + if type(semantics) is not SkillPlanSemantics: + raise TypeError("_trajectory_semantics must return SkillPlanSemantics") + plan = self.build_plan( + request, + context, + success=owned.planning_success, + trajectory=timed, + expected_effects=semantics.expected_effects, + effect_candidates=semantics.effect_candidates, + effect_verification=semantics.effect_verification, + replannable=owned.replannable, + diagnostics=owned.diagnostics, + segment_lengths=( + {s.name: s.waypoint_count for s in owned.segments} + if owned.segments + else None + ), + scene_dependency_monitor_until=semantics.scene_dependency_monitor_until, + scene_dependency_end_segment=semantics.scene_dependency_end_segment, + ) + return replace( + plan, + trajectory_input=replace( + owned, + trajectory=plan.joint_trajectory, + planning_success=plan.plan_success, + ), + ) + + def _trajectory_semantics( + self, + request: ResolvedActionRequest[GoalT, OptionsT], + context: PlanningContext, + motion: SkillTrajectoryInput, + ) -> SkillPlanSemantics: + """Supply Skill-owned effects without planning, executing or committing state.""" + raise NotImplementedError( + f"Skill {self.skill_id!r} has no trajectory semantics" + ) + def build_plan( self, request: ResolvedActionRequest[GoalT, OptionsT], diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index 258c41aa0..85b37ffef 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -19,21 +19,24 @@ from __future__ import annotations from types import MappingProxyType +from dataclasses import replace from typing import Callable, Iterable, Mapping, TYPE_CHECKING import torch +from embodichain.compute.trajectory import differentiate_positions from .affordance_sampling import AffordanceSamplingContext -from .bindings import ActionBinding +from .bindings import ActionBinding, JointPositionTarget from .core import AtomicAction, SkillDescriptor from .control import ActionControlOverrides, ControlPartCommandProfile from .invocation import ActionInvocation, GoalT, OptionsT, ResolvedActionRequest -from .plans import ActionPlan, CompiledTrajectory, TimedTrajectory +from .plans import ActionPlan, CompiledTrajectory, TimedTrajectory, PlannerDiagnostics from .policies import MotionPolicy, RecoveryPolicy from .runtime import ActionPlanningServices from .scene import SceneProvider from .state import PlanningContext, RobotObservation, SceneSnapshot, TaskState from .tracking import TrackingPolicy, TrackingRuntime +from .trajectory_assembly import SkillTrajectoryInput if TYPE_CHECKING: from embodichain.lab.sim.objects import Robot @@ -49,6 +52,10 @@ ] +NominalPlanProvider = Callable[[ResolvedActionRequest, PlanningContext], ActionPlan] +"""Supply a complete nominal plan without invoking ordinary action planning.""" + + class AtomicActionEngine: """Own planning resources and coordinate side-effect-free atomic actions.""" @@ -287,6 +294,13 @@ def register(self, action: AtomicAction, *, replace: bool = False) -> None: if not isinstance(action, AtomicAction): raise TypeError("action must be an AtomicAction instance.") descriptor = action.descriptor() + if ( + descriptor.trajectory_contract is not None + and type(action)._trajectory_semantics is AtomicAction._trajectory_semantics + ): + raise ValueError( + "A trajectory contract requires a Skill semantic assembler" + ) existing = self._actions.get(descriptor.skill_id) if existing is not None and existing is not action and not replace: raise ValueError( @@ -324,12 +338,15 @@ def plan_request( context: PlanningContext | None = None, *, plan_transform: PlanTransform | None = None, + nominal_plan_provider: NominalPlanProvider | None = None, ) -> ActionPlan: """Plan an already-resolved request without rebuilding its snapshot. Args: request: Immutable request previously returned by :meth:`resolve`. context: Optional latest planning state; captured when omitted. + nominal_plan_provider: Optional complete-plan source. When supplied, + ordinary action planning is bypassed; failures never fall back. plan_transform: Optional callback scoped to this planning call. The callback receives and must return independently owned plan values. @@ -340,6 +357,7 @@ def plan_request( request, context, plan_transform=plan_transform, + nominal_plan_provider=nominal_plan_provider, ) def _resolve( @@ -374,6 +392,7 @@ def _plan_request( context: PlanningContext | None = None, *, plan_transform: PlanTransform | None = None, + nominal_plan_provider: NominalPlanProvider | None = None, ) -> ActionPlan: """Plan an already-resolved request without rebuilding its snapshot. @@ -384,6 +403,8 @@ def _plan_request( Args: request: Immutable request previously returned by :meth:`_resolve`. context: Optional latest planning state; captured when omitted. + nominal_plan_provider: Optional complete-plan source. When supplied, + ordinary action planning is bypassed; failures never fall back. Returns: Validated side-effect-free action plan. @@ -397,7 +418,15 @@ def _plan_request( ) current = self.initial_context() if context is None else context self._validate_context(current) - plan = action.plan(request, current) + if nominal_plan_provider is None: + plan = action.plan(request, current) + else: + if not callable(nominal_plan_provider): + raise TypeError("nominal_plan_provider must be callable or None.") + supplied = nominal_plan_provider(request, current) + if not isinstance(supplied, ActionPlan): + raise TypeError("nominal_plan_provider must return an ActionPlan.") + plan = supplied.snapshot() self._validate_plan(plan, current, request) if plan_transform is not None: if not callable(plan_transform): @@ -416,12 +445,15 @@ def plan( context: PlanningContext | None = None, *, plan_transform: PlanTransform | None = None, + nominal_plan_provider: NominalPlanProvider | None = None, ) -> ActionPlan: """Plan one registered invocation through the engine-owned backend. Args: invocation: Grounded request for a registered skill. context: Optional latest planning state; captured when omitted. + nominal_plan_provider: Optional complete-plan source. When supplied, + ordinary action planning is bypassed; failures never fall back. Returns: Validated action plan. @@ -431,7 +463,12 @@ def plan( """ current = self.initial_context() if context is None else context request = self._resolve(invocation) - return self._plan_request(request, current, plan_transform=plan_transform) + return self._plan_request( + request, + current, + plan_transform=plan_transform, + nominal_plan_provider=nominal_plan_provider, + ) def initial_context( self, @@ -573,6 +610,7 @@ def start( *, eligible_mask: torch.Tensor | None = None, plan_transform: PlanTransform | None = None, + nominal_plan_provider: NominalPlanProvider | None = None, ) -> ExecutionSession: """Start incremental execution for a grounded invocation sequence. @@ -583,6 +621,8 @@ 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. + nominal_plan_provider: Optional source retained for initial planning, + invocation revisions and recovery throughout the session. plan_transform: Optional transform scoped to every planning call in the returned execution session. @@ -598,6 +638,7 @@ def start( initial, eligible_mask=eligible_mask, plan_transform=plan_transform, + nominal_plan_provider=nominal_plan_provider, ) def _validate_context(self, context: PlanningContext) -> None: @@ -666,6 +707,184 @@ def _validate_plan( "Action plan must record the planning collision-world revision." ) + def build_plan_from_trajectory( + self, + request: ResolvedActionRequest, + context: PlanningContext, + motion: SkillTrajectoryInput, + ) -> ActionPlan: + """Admit external motion and delegate complete assembly to its Skill. + + Args: + request: Grounded invocation owning the trajectory's endpoint claims. + context: Observed state restored to this motion's initial conditions. + motion: Source-independent motion with exact stages and event indices. + + Returns: + Validated, owned plan without invoking ordinary planning. + + This checks discrete motion limits and authorization, not collision + freedom or measured task success. Source/host validation retains those + responsibilities and registered source recovery admission is unchanged. + """ + if not isinstance(request, ResolvedActionRequest): + raise TypeError("request must be a ResolvedActionRequest") + self._validate_context(context) + action = self._actions.get(request.skill_id) + if action is None: + raise KeyError( + f"No atomic action registered for skill {request.skill_id!r}." + ) + contract = action.descriptor().trajectory_contract + if contract is None: + raise ValueError( + f"Skill {request.skill_id!r} does not support trajectory assembly" + ) + if type(motion) is not SkillTrajectoryInput: + raise TypeError("motion must be exactly SkillTrajectoryInput") + owned = motion.snapshot() + contract.validate_input(owned) + self._validate_external_trajectory(request, context, owned) + plan = action.assemble_trajectory_plan(request, context, owned) + self._validate_plan(plan, context, request) + return plan.snapshot() + + def _validate_external_trajectory( + self, + request: ResolvedActionRequest, + context: PlanningContext, + motion: SkillTrajectoryInput, + ) -> None: + """Check restored initial state, authorized joints and discrete limits.""" + timed = motion.trajectory + q = timed.positions + if q.shape[0] != context.batch_size or q.shape[2] != self.robot.dof: + raise ValueError("Trajectory dimensions must match the planning context") + if q.device != self.device or not torch.equal(timed.env_ids, context.env_ids): + raise ValueError("Trajectory must share context device and ordered env_ids") + if q.dtype != context.robot.qpos.dtype: + raise ValueError("Trajectory and planning context must share a joint dtype") + active = motion.planning_success + if not torch.allclose( + q[active, 0], context.robot.qpos[active], atol=1e-4, rtol=0 + ): + raise ValueError("Restore the exact trajectory initial state before replay") + if not motion.replannable and ( + request.revision + or request.recovery_policy.max_replans + or request.recovery_policy.max_action_retries + ): + raise ValueError( + "Fixed trajectory replay cannot replan or retry without a fresh reset" + ) + limits = self.robot.get_qpos_limits(env_ids=context.env_ids) + if not isinstance(limits, torch.Tensor) or limits.shape != ( + context.batch_size, + self.robot.dof, + 2, + ): + raise ValueError("Live robot must expose full-joint position limits") + limits = limits.to(q) + if not torch.isfinite(limits).all() or (limits[..., 0] > limits[..., 1]).any(): + raise ValueError("Live joint position limits must be finite and ordered") + if ( + (q[active] < limits[active, None, :, 0] - 1e-6) + | (q[active] > limits[active, None, :, 1] + 1e-6) + ).any(): + raise ValueError("Trajectory violates live robot joint position limits") + velocity_limits = self.robot.get_qvel_limits(env_ids=context.env_ids) + if ( + not isinstance(velocity_limits, torch.Tensor) + or velocity_limits.shape != (context.batch_size, self.robot.dof) + or not torch.isfinite(velocity_limits).all() + or (velocity_limits <= 0).any() + ): + raise ValueError( + "Live robot must expose positive finite joint velocity limits" + ) + velocities = timed.velocities + if velocities is None: + velocities = differentiate_positions(q, timed.dt) + if ( + velocities[active].abs() > velocity_limits.to(q)[active, None] + 1e-6 + ).any(): + raise ValueError("Trajectory violates live robot joint velocity limits") + intervals = timed.dt[:, 1:] + if (intervals[active] <= 0).any(): + raise ValueError( + "External trajectory arrival intervals must be positive after its first sample" + ) + speeds = q[active].diff(dim=1).abs() / intervals[active, :, None] + if (speeds > velocity_limits.to(q)[active, None] + 1e-6).any(): + raise ValueError( + "Trajectory interval speeds violate live robot joint velocity limits" + ) + claimed = { + joint + for endpoint in request.binding.endpoints + if isinstance(endpoint.target, JointPositionTarget) + for joint in endpoint.target.joint_ids + } + mimics = getattr(self.robot, "mimic_ids", ()) + if isinstance(mimics, (list, tuple, set, torch.Tensor)): + claimed.update(int(index) for index in mimics) + unclaimed = sorted(set(range(self.robot.dof)) - claimed) + if unclaimed and not torch.allclose( + q[active][:, :, unclaimed], + context.robot.qpos[active][:, None, unclaimed].expand( + -1, timed.waypoint_count, -1 + ), + atol=1e-6, + rtol=0, + ): + raise ValueError( + "Trajectory moves joints outside the call's authorized endpoints" + ) + + def build_plan_from_recorded_trajectory( + self, + request: ResolvedActionRequest, + context: PlanningContext, + trajectory: TimedTrajectory, + *, + segment_lengths: dict[str, int], + diagnostics: PlannerDiagnostics, + ) -> ActionPlan: + """Adapt the original recorded-path API to generic trajectory assembly. + + Args: + request: Resolved invocation selected by the Semantic Call. + context: Live planning state after restoring the recording start. + trajectory: Source-normalized full-joint path on the host clock. + segment_lengths: Skill-owned semantic stage mapping. + diagnostics: Source evidence and normalization lineage. + + Returns: + Validated complete nominal plan, without running ordinary planning. + """ + action = self._actions.get(request.skill_id) + contract = None if action is None else action.descriptor().trajectory_contract + if contract is None: + raise ValueError( + f"Skill {request.skill_id!r} does not support recorded plan assembly" + ) + segments = action._build_segments( + segment_lengths, frame_count=trajectory.waypoint_count + ) + landmarks = contract.stage_landmarks(segments, trajectory.waypoint_count) + return self.build_plan_from_trajectory( + request, + context, + SkillTrajectoryInput( + schema_id=contract.schema_id, + schema_revision=contract.revision, + trajectory=trajectory, + segments=segments, + landmarks=landmarks, + diagnostics=diagnostics, + ), + ) + def rebuild_plan_from_trajectory( self, request: ResolvedActionRequest, @@ -713,6 +932,23 @@ def rebuild_plan_from_trajectory( original.dt, ): raise ValueError("same-grid plan rebuilding requires identical dt.") + if owned_plan.trajectory_input is not None: + source = owned_plan.trajectory_input + if source.geometry is not None: + events = list(source.landmarks.values()) + if not torch.equal( + trajectory.positions[:, events], original.positions[:, events] + ): + raise ValueError( + "Changed geometric event poses require a new trajectory assembly input" + ) + return self.build_plan_from_trajectory( + request, + context, + replace( + source, trajectory=trajectory, diagnostics=owned_plan.diagnostics + ), + ) action = self._actions.get(request.skill_id) if action is None: raise KeyError( diff --git a/embodichain/lab/sim/atomic_actions/execution.py b/embodichain/lab/sim/atomic_actions/execution.py index 9f967c8fc..db2332adb 100644 --- a/embodichain/lab/sim/atomic_actions/execution.py +++ b/embodichain/lab/sim/atomic_actions/execution.py @@ -56,7 +56,7 @@ ) if TYPE_CHECKING: - from .engine import AtomicActionEngine, PlanTransform + from .engine import AtomicActionEngine, PlanTransform, NominalPlanProvider class ExecutionStatus(str, Enum): @@ -326,6 +326,7 @@ def __init__( *, eligible_mask: torch.Tensor | None = None, plan_transform: PlanTransform | None = None, + nominal_plan_provider: NominalPlanProvider | None = None, ) -> None: if not invocations: raise ValueError("ExecutionSession requires at least one invocation.") @@ -334,6 +335,9 @@ def __init__( raise TypeError("plan_transform must be callable or None.") self._engine = engine self._plan_transform = plan_transform + if nominal_plan_provider is not None and not callable(nominal_plan_provider): + raise TypeError("nominal_plan_provider must be callable or None.") + self._nominal_plan_provider = nominal_plan_provider self._requests: tuple[ResolvedActionRequest, ...] = tuple( engine._resolve(invocation) for invocation in invocations ) @@ -612,6 +616,7 @@ def _install_prepared_revision( replacement, replacement_context, plan_transform=self._plan_transform, + nominal_plan_provider=self._nominal_plan_provider, ) self._validate_destination_continuity( replacement_plan, @@ -1157,6 +1162,7 @@ def _plan_current( request, context, plan_transform=self._plan_transform, + nominal_plan_provider=self._nominal_plan_provider, ) self._install_plan(plan, context, event_kind) diff --git a/embodichain/lab/sim/atomic_actions/plans.py b/embodichain/lab/sim/atomic_actions/plans.py index f29493727..cfaf968ff 100644 --- a/embodichain/lab/sim/atomic_actions/plans.py +++ b/embodichain/lab/sim/atomic_actions/plans.py @@ -22,7 +22,7 @@ from copy import deepcopy from dataclasses import dataclass, field from types import MappingProxyType -from typing import Any, Mapping, Sequence +from typing import Any, Mapping, Sequence, TYPE_CHECKING import torch @@ -38,6 +38,9 @@ TrackingPolicy, ) +if TYPE_CHECKING: + from .trajectory_assembly import SkillTrajectoryInput + def _validate_optional_trajectory_field( value: torch.Tensor | None, @@ -528,6 +531,8 @@ class ActionPlan: effect_verification: EffectVerificationRequirement | None = None invocation_id: str | None = None invocation_revision: int = 0 + trajectory_input: SkillTrajectoryInput | None = None + """Owned semantic assembly input for rebuilding a supported joint plan.""" def __post_init__(self) -> None: if not isinstance(self.skill_id, str) or not self.skill_id: @@ -814,6 +819,34 @@ def __post_init__(self) -> None: raise ValueError( "scene_dependency_end_segment requires scene_dependencies." ) + if self.trajectory_input is not None: + from .trajectory_assembly import SkillTrajectoryInput + + if type(self.trajectory_input) is not SkillTrajectoryInput: + raise TypeError("trajectory_input must be SkillTrajectoryInput or None") + source = self.trajectory_input.snapshot() + timed = self.joint_trajectory + if timed is None or any( + not torch.equal(getattr(source.trajectory, name), getattr(timed, name)) + for name in ("positions", "dt", "env_ids") + ): + raise ValueError( + "Assembly input must describe the plan's joint trajectory" + ) + source_segments = source.segments or ( + TrajectorySegment(self.skill_id, 0, source.trajectory.waypoint_count), + ) + if source_segments != segments or not torch.equal( + source.planning_success, self.plan_success + ): + raise ValueError( + "Assembly input stages and planning mask must match the plan" + ) + if source.replannable != self.replannable: + raise ValueError( + "Assembly input must preserve the plan's recovery capability" + ) + object.__setattr__(self, "trajectory_input", source) object.__setattr__(self, "plan_success", self.plan_success.clone()) object.__setattr__(self, "commands", self.commands.snapshot()) object.__setattr__( @@ -907,6 +940,7 @@ def snapshot(self) -> ActionPlan: effect_verification=self.effect_verification, invocation_id=self.invocation_id, invocation_revision=self.invocation_revision, + trajectory_input=self.trajectory_input, ) @property diff --git a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py index 58bd3a33d..cc4375375 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/move_joints.py +++ b/embodichain/lab/sim/atomic_actions/primitives/move_joints.py @@ -30,6 +30,11 @@ ResolvedActionRequest, ) from embodichain.lab.sim.atomic_actions.plans import ActionPlan +from embodichain.lab.sim.atomic_actions.trajectory_assembly import ( + SkillPlanSemantics, + SkillTrajectoryContract, + SkillTrajectoryInput, +) from embodichain.lab.sim.atomic_actions.requirements import ( JOINT_POSITION_CAPABILITY, SkillBindingContract, @@ -83,6 +88,12 @@ class MoveJoints(AtomicAction[JointPositionGoal, MoveJointsOptions]): GoalType: ClassVar[type] = JointPositionGoal OptionsType: ClassVar[type] = MoveJointsOptions agent_visible: ClassVar[bool] = True + trajectory_contract: ClassVar[SkillTrajectoryContract] = SkillTrajectoryContract( + schema_id="move_joints.qpos", + revision="1", + landmark_names=("goal",), + landmark_stages={"goal": None}, + ) binding_contract: ClassVar[SkillBindingContract] = SkillBindingContract( slots=( make_motion_slot( @@ -131,12 +142,92 @@ def _plan( env_ids=context.env_ids, control_dt=context.require_control_dt(), ) - return self.build_plan( + return self.assemble_trajectory_plan( request, context, - success=success, - trajectory=trajectory, + SkillTrajectoryInput( + schema_id=self.trajectory_contract.schema_id, + schema_revision=self.trajectory_contract.revision, + trajectory=trajectory, + landmarks={"goal": trajectory.waypoint_count - 1}, + planning_success=success, + replannable=True, + ), + ) + + def _trajectory_semantics( + self, + request: ResolvedActionRequest[JointPositionGoal, MoveJointsOptions], + context: PlanningContext, + motion: SkillTrajectoryInput, + ) -> SkillPlanSemantics: + """Check ordered joint goals without introducing symbolic task effects.""" + target = request.binding.endpoint("primary", "motion").require_target( + JointPositionTarget + ) + goals = resolve_joint_target( + self._resolve_target_qpos(request.goal, request=request, context=context), + num_envs=context.batch_size, + joint_dof=len(target.joint_ids), + control_part=target.control_part, + device=self.device, ) + if goals.ndim == 2: + goals = goals[:, None] + positions = motion.trajectory.positions[:, :, list(target.joint_ids)] + if motion.landmarks["goal"] != motion.trajectory.waypoint_count - 1: + raise ValueError("MoveJoints goal event must be the terminal sample") + # This checks the reference path, not measured controller acceptance. + tolerance = 1e-4 + for row in motion.planning_success.nonzero().flatten().tolist(): + if not torch.allclose( + positions[row, -1], goals[row, -1].to(positions), atol=tolerance, rtol=0 + ): + raise ValueError( + "MoveJoints trajectory does not reach its requested joint goal" + ) + path = positions[row] + if len(path) == 1: + if not torch.allclose( + goals[row].to(path), + path.expand_as(goals[row]), + atol=tolerance, + rtol=0, + ): + raise ValueError( + "MoveJoints trajectory must visit requested waypoints in order" + ) + continue + edges = path[1:] - path[:-1] + squared_lengths = edges.square().sum(dim=-1) + cursor = 0.0 + for waypoint in goals[row]: + # Control-grid retiming can put an interior waypoint between + # samples. Check the ordered reference polyline, not only exact + # sample equality; this still does not certify physical motion. + delta = waypoint.to(path) - path[:-1] + fractions = ( + (delta * edges).sum(dim=-1) + / squared_lengths.clamp_min(torch.finfo(path.dtype).tiny) + ).clamp(0, 1) + closest = path[:-1] + fractions[:, None] * edges + progress = torch.arange( + len(edges), device=path.device, dtype=torch.float64 + ) + fractions.to(torch.float64) + matches = ( + ( + ((closest - waypoint.to(path)).abs() <= tolerance).all(dim=-1) + & (progress >= cursor - 1e-8) + ) + .nonzero() + .flatten() + ) + if not len(matches): + raise ValueError( + "MoveJoints trajectory must visit requested waypoints in order" + ) + cursor = float(progress[matches[0]]) + return SkillPlanSemantics() def _resolve_target_qpos( self, diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index eb02b729d..2c5dd3ef4 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -68,6 +68,11 @@ TimedTrajectory, normalize_success_mask, ) +from embodichain.lab.sim.atomic_actions.trajectory_assembly import ( + SkillPlanSemantics, + SkillTrajectoryContract, + SkillTrajectoryInput, +) from embodichain.lab.sim.atomic_actions.policies import MotionPolicy from embodichain.lab.sim.atomic_actions.requirements import ( BATCH_INVERSE_KINEMATICS_CAPABILITY, @@ -288,12 +293,55 @@ def __post_init__(self) -> None: ) +@dataclass(frozen=True, slots=True, eq=False) +class PickTrajectoryGeometry: + """Expected grasp pose selected by ordinary Pick planning. + + Args: + grasp_pose: Batched expected grasp transforms, shape ``(B, 4, 4)``. + + This is planning geometry, not independently measured grasp evidence. + Imported joint paths may omit it so PickUp derives the pose by FK at the + explicit grasp event instead. + """ + + grasp_pose: torch.Tensor + + def __post_init__(self) -> None: + pose = self.grasp_pose + if ( + not isinstance(pose, torch.Tensor) + or pose.ndim != 3 + or pose.shape[1:] != (4, 4) + ): + raise ValueError("grasp_pose must have shape (B, 4, 4)") + if not torch.isfinite(pose).all(): + raise ValueError("grasp_pose must contain finite transforms") + object.__setattr__(self, "grasp_pose", pose.detach().clone()) + + def snapshot(self) -> PickTrajectoryGeometry: + """Return independently owned expected grasp geometry. + + Returns: + A distinct value with copied grasp transforms. + """ + return PickTrajectoryGeometry(self.grasp_pose) + + class PickUp(AtomicAction[GraspGoal, PickUpOptions]): """Approach a grasp pose, close the gripper, lift.""" skill_id: ClassVar[str] = "pick_up" GoalType: ClassVar[type] = GraspGoal OptionsType: ClassVar[type] = PickUpOptions + trajectory_contract: ClassVar[SkillTrajectoryContract] = SkillTrajectoryContract( + schema_id="pick_up.qpos", + revision="1", + segment_names=("approach", "close", "lift"), + landmark_names=("grasp",), + landmark_stages={"grasp": "close"}, + geometry_type=PickTrajectoryGeometry, + ) binding_contract: ClassVar[SkillBindingContract] = SkillBindingContract( slots=( make_manipulation_slot( @@ -584,52 +632,116 @@ def _plan( name="Pick-up trajectory success", ) - object_to_eef = torch.bmm(pose_inv(object_pose), grasp_xpos) - held = HeldObjectState( - semantics=sem, object_to_eef=object_to_eef, grasp_xpos=grasp_xpos - ) - coordinated_updates = { - key: None for key in state.coordinated_held_objects if task_state_key in key - } - if target.object_pose is None: - monitored_object_id = sem.entity_id - elif isinstance(target.object_pose, SceneEntityPose): - monitored_object_id = target.object_pose.entity_id - else: - monitored_object_id = None - return self.build_plan( + segments = self._build_segments(segment_lengths, frame_count=full.shape[1]) + return self.assemble_trajectory_plan( request, context, - success=success_mask, - trajectory=TimedTrajectory.from_uniform_step( - full, - env_ids=context.env_ids, - step_dt=context.require_control_dt(), + SkillTrajectoryInput( + schema_id=self.trajectory_contract.schema_id, + schema_revision=self.trajectory_contract.revision, + trajectory=TimedTrajectory.from_uniform_step( + full, + env_ids=context.env_ids, + step_dt=context.require_control_dt(), + ), + segments=segments, + landmarks={ + "grasp": segment_lengths["approach"] + segment_lengths["close"] - 1 + }, + planning_success=success_mask, + geometry=PickTrajectoryGeometry(grasp_xpos), + diagnostics=PlannerDiagnostics( + backend=self.planning_services.planner_name, + metadata=( + {} + if affordance_sample is None + else { + "affordance_sample": {"grasp": affordance_sample.metadata} + } + ), + ), + replannable=True, ), + ) + + def _trajectory_semantics( + self, + request: ResolvedActionRequest[GraspGoal, PickUpOptions], + context: PlanningContext, + motion: SkillTrajectoryInput, + ) -> SkillPlanSemantics: + """Build the same expected grasp effects for planned and imported motion.""" + target = self.require_goal(request) + approach, close, _ = motion.segments + if motion.landmarks["grasp"] != close.stop - 1: + raise ValueError("PickUp grasp event must be the final close sample") + arm = request.binding.endpoint("primary", "motion") + grasp = request.binding.endpoint("primary", "grasp") + manipulator = arm.require_target(JointPositionTarget) + key = require_shared_task_state_key( + arm, grasp, participant="PickUp primary participant" + ) + geometry = motion.geometry + if geometry is None: + grasp_pose = self.robot.compute_fk( + motion.trajectory.positions[ + :, motion.landmarks["grasp"], list(manipulator.joint_ids) + ], + name=manipulator.control_part, + env_ids=context.env_ids, + to_matrix=True, + ) + else: + assert isinstance(geometry, PickTrajectoryGeometry) + grasp_pose = geometry.grasp_pose + if ( + grasp_pose.shape != (context.batch_size, 4, 4) + or grasp_pose.device != self.device + ): + raise ValueError( + "Pick grasp geometry must match the planning batch and device" + ) + object_pose = ( + _resolve_object_pose(target.semantics, context, name="pickup_object_pose") + if target.object_pose is None + else resolve_pose_target( + resolve_pose_goal(target.object_pose, context, name="object_pose"), + num_envs=context.batch_size, + device=self.device, + ) + ) + held = HeldObjectState( + semantics=target.semantics, + object_to_eef=torch.bmm(pose_inv(object_pose), grasp_pose), + grasp_xpos=grasp_pose, + ) + monitored = ( + target.semantics.entity_id + if target.object_pose is None + else ( + target.object_pose.entity_id + if isinstance(target.object_pose, SceneEntityPose) + else None + ) + ) + return SkillPlanSemantics( expected_effects=StateDelta( - held_object_updates={task_state_key: held}, - coordinated_held_object_updates=coordinated_updates, - ), - diagnostics=PlannerDiagnostics( - backend=self.planning_services.planner_name, - metadata=( - {} - if affordance_sample is None - else {"affordance_sample": {"grasp": affordance_sample.metadata}} - ), + held_object_updates={key: held}, + coordinated_held_object_updates={ + group: None + for group in context.coordinated_held_objects + if key in group + }, ), - segment_lengths=segment_lengths, - # Once the approach is dispatched, contact can move the acquired - # object. Keep downstream late-bound targets monitored so a changed - # destination is resolved again before the action continues. scene_dependency_monitor_until=( {} - if monitored_object_id is None + if monitored is None else { - monitored_object_id: max( + monitored: max( 1, math.ceil( - segment_lengths["approach"] * options.grasp_commit_fraction + approach.waypoint_count + * request.skill_options.grasp_commit_fraction ), ) } @@ -901,4 +1013,4 @@ def _upright_adjusted_grasp_poses( return adjusted_grasp_xpos -__all__ = ["GraspGoal", "PickUp", "PickUpOptions"] +__all__ = ["GraspGoal", "PickUp", "PickUpOptions", "PickTrajectoryGeometry"] diff --git a/embodichain/lab/sim/atomic_actions/trajectory_assembly.py b/embodichain/lab/sim/atomic_actions/trajectory_assembly.py new file mode 100644 index 000000000..c577d24ad --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/trajectory_assembly.py @@ -0,0 +1,346 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Source-independent inputs and Skill semantics for joint-plan assembly.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Protocol, runtime_checkable + +import torch + +from .effects import StateDelta +from .plans import ( + EffectVerificationRequirement, + PlannerDiagnostics, + TimedTrajectory, + TrajectorySegment, + normalize_success_mask, +) + +__all__ = [ + "SkillTrajectoryGeometry", + "SkillTrajectoryContract", + "SkillTrajectoryInput", + "SkillPlanSemantics", +] + + +def _name(value: str) -> None: + if type(value) is not str or not value or value.strip() != value: + raise ValueError("Trajectory schema names must be nonempty stripped strings") + + +def _names(values: tuple[str, ...]) -> tuple[str, ...]: + result = tuple(values) + for value in result: + _name(value) + if len(set(result)) != len(result): + raise ValueError("Trajectory schema names must be unique") + return result + + +@runtime_checkable +class SkillTrajectoryGeometry(Protocol): + """Skill-owned geometric facts with an independently owned snapshot.""" + + def snapshot(self) -> SkillTrajectoryGeometry: + """Return an independently owned geometric-facts value. + + Returns: + A distinct value of the same type with owned tensor storage. + """ + + +@dataclass(frozen=True, slots=True) +class SkillTrajectoryContract: + """Versioned, declarative opt-in to trajectory-backed plan assembly. + + Args: + schema_id: Stable Skill-specific trajectory schema identity. + revision: Exact schema version required by an input. + segment_names: Ordered semantic stages; empty uses the framework's + default Skill-named motion stage. + landmark_names: Exact required event IDs on the execution grid. + landmark_stages: Optional event-to-stage-end constraints. ``None`` + refers to the terminal sample of the full path. + geometry_type: Optional concrete geometric-facts type accepted by the Skill. + Omitting geometry asks the Skill to derive facts from the trajectory + and current context; it never authorizes arbitrary metadata. + supports_batch: Whether rows sharing one stage layout are supported. + + This contract does not grant expansion, source admission or recovery + capabilities. Changing execution grids requires an explicitly mapped new + input; the existing same-grid rebuild API retains its timing restriction. + """ + + schema_id: str + revision: str + segment_names: tuple[str, ...] = () + landmark_names: tuple[str, ...] = () + landmark_stages: Mapping[str, str | None] = field(default_factory=dict) + geometry_type: type[SkillTrajectoryGeometry] | None = None + supports_batch: bool = True + + def __post_init__(self) -> None: + _name(self.schema_id) + _name(self.revision) + object.__setattr__(self, "segment_names", _names(self.segment_names)) + object.__setattr__(self, "landmark_names", _names(self.landmark_names)) + stages = dict(self.landmark_stages) + if set(stages) - set(self.landmark_names) or any( + value is not None and value not in self.segment_names + for value in stages.values() + ): + raise ValueError( + "Landmark stage constraints must name declared events and stages" + ) + object.__setattr__(self, "landmark_stages", MappingProxyType(stages)) + if self.geometry_type is not None and ( + not isinstance(self.geometry_type, type) + or not issubclass(self.geometry_type, SkillTrajectoryGeometry) + ): + raise TypeError("geometry_type must implement SkillTrajectoryGeometry") + if type(self.supports_batch) is not bool: + raise TypeError("supports_batch must be a bool") + + def validate_input(self, motion: SkillTrajectoryInput) -> None: + """Check an owned input against this exact semantic schema. + + Args: + motion: Trajectory and semantic annotations to admit. + """ + if type(motion) is not SkillTrajectoryInput: + raise TypeError("motion must be exactly SkillTrajectoryInput") + if (motion.schema_id, motion.schema_revision) != ( + self.schema_id, + self.revision, + ): + raise ValueError( + "Trajectory schema identity or revision does not match Skill" + ) + if tuple(s.name for s in motion.segments) != self.segment_names: + raise ValueError("Trajectory semantic stages do not match Skill contract") + if set(motion.landmarks) != set(self.landmark_names): + raise ValueError("Trajectory event landmarks do not match Skill contract") + if any( + motion.landmarks[name] != index + for name, index in self.stage_landmarks( + motion.segments, motion.trajectory.waypoint_count + ).items() + ): + raise ValueError( + "Trajectory event landmark violates its declared stage boundary" + ) + if ( + motion.geometry is not None + and type(motion.geometry) is not self.geometry_type + ): + raise TypeError("Trajectory geometric facts do not match Skill contract") + if motion.trajectory.batch_size > 1 and not self.supports_batch: + raise ValueError("Skill trajectory contract requires B=1") + + def stage_landmarks( + self, + segments: tuple[TrajectorySegment, ...], + waypoint_count: int, + ) -> Mapping[str, int]: + """Resolve declared stage-end events on a supplied execution grid. + + Args: + segments: Semantic ranges in this version of the trajectory. + waypoint_count: Number of valid trajectory samples. + + Returns: + Bound event indices; interior events still require explicit annotation. + """ + if type(waypoint_count) is not int or waypoint_count < 1: + raise ValueError("Event resolution requires a positive sample count") + by_name = {segment.name: segment for segment in segments} + if any( + stage is not None and stage not in by_name + for stage in self.landmark_stages.values() + ): + raise ValueError("Event resolution is missing a declared semantic stage") + return MappingProxyType( + { + name: waypoint_count - 1 if stage is None else by_name[stage].stop - 1 + for name, stage in self.landmark_stages.items() + } + ) + + +@dataclass(frozen=True, slots=True, eq=False) +class SkillTrajectoryInput: + """Owned, source-neutral motion supplied to a Skill's semantic assembler. + + Args: + schema_id: Selected Skill trajectory contract identity. + schema_revision: Exact selected contract version. + trajectory: Full-joint motion on the host's execution clock. + segments: Contiguous half-open semantic ranges covering the path, or + empty for the framework's default Skill-named motion stage. + landmarks: Stable event IDs mapped to actual samples, never exclusive + end boundaries. A stage's last sample is ``stop - 1``. + planning_success: Per-row planning eligibility, not physical success. + geometry: Optional typed Skill facts; data decoding belongs to the source. + diagnostics: Source/planner provenance, without executable callbacks. + replannable: Whether the supplying planning path can regenerate this + input. Registered source admission still constrains recovery policy. + + Semantic stages do not grant motion-expansion permissions. Existing + ``TrajectoryPhase`` values remain owned by the augmentation boundary. + """ + + schema_id: str + schema_revision: str + trajectory: TimedTrajectory + segments: tuple[TrajectorySegment, ...] = () + landmarks: Mapping[str, int] = field(default_factory=dict) + planning_success: bool | torch.Tensor = True + geometry: SkillTrajectoryGeometry | None = None + diagnostics: PlannerDiagnostics | None = None + replannable: bool = False + + def __post_init__(self) -> None: + _name(self.schema_id) + _name(self.schema_revision) + if not isinstance(self.trajectory, TimedTrajectory): + raise TypeError("trajectory must be a TimedTrajectory") + trajectory = self.trajectory.snapshot() + if not trajectory.waypoint_count: + raise ValueError("Trajectory assembly requires at least one sample") + segments = tuple(self.segments) + offset = 0 + for segment in segments: + if type(segment) is not TrajectorySegment: + raise TypeError("segments must contain TrajectorySegment values") + if segment.start != offset: + raise ValueError("Semantic stages must be contiguous and start at zero") + offset = segment.stop + _names(tuple(s.name for s in segments)) + if segments and offset != trajectory.waypoint_count: + raise ValueError("Semantic stages must cover the complete trajectory") + landmarks = dict(self.landmarks) + for name, index in landmarks.items(): + _name(name) + if type(index) is not int or not 0 <= index < trajectory.waypoint_count: + raise ValueError("Event landmarks must index actual trajectory samples") + mask = normalize_success_mask( + self.planning_success, + num_envs=trajectory.batch_size, + device=trajectory.positions.device, + name="Trajectory planning success", + ) + geometry = self.geometry + if geometry is not None: + if not isinstance(geometry, SkillTrajectoryGeometry): + raise TypeError("geometry must implement SkillTrajectoryGeometry") + copied = geometry.snapshot() + if type(copied) is not type(geometry) or copied is geometry: + raise TypeError( + "Geometric snapshots must be distinct values of the same type" + ) + geometry = copied + diagnostics = self.diagnostics + if diagnostics is not None: + if type(diagnostics) is not PlannerDiagnostics: + raise TypeError("diagnostics must be PlannerDiagnostics or None") + diagnostics = PlannerDiagnostics( + diagnostics.backend, + diagnostics.messages, + diagnostics.metadata, + diagnostics.failure, + ) + if type(self.replannable) is not bool: + raise TypeError("replannable must be a bool") + object.__setattr__(self, "trajectory", trajectory) + object.__setattr__(self, "segments", segments) + object.__setattr__(self, "landmarks", MappingProxyType(landmarks)) + object.__setattr__(self, "planning_success", mask.clone()) + object.__setattr__(self, "geometry", geometry) + object.__setattr__(self, "diagnostics", diagnostics) + + def snapshot(self) -> SkillTrajectoryInput: + """Return an independently owned input for a planning boundary. + + Returns: + Revalidated trajectory, annotations and geometric facts. + """ + return SkillTrajectoryInput( + schema_id=self.schema_id, + schema_revision=self.schema_revision, + trajectory=self.trajectory, + segments=self.segments, + landmarks=self.landmarks, + planning_success=self.planning_success, + geometry=self.geometry, + diagnostics=self.diagnostics, + replannable=self.replannable, + ) + + +@dataclass(frozen=True, slots=True, eq=False) +class SkillPlanSemantics: + """Uncommitted Skill effects and scene-monitor boundaries for one motion. + + Args: + expected_effects: Terminal symbolic effects handled by the existing runtime. + effect_candidates: Intermediate attachment baselines for gates and guards. + effect_verification: Optional explicit physical-effect requirement. + scene_dependency_monitor_until: Exclusive execution-index monitor bounds, + computed from this input's current events and stages. + scene_dependency_end_segment: Optional named last monitored stage. + + Constructing this value never changes task state or establishes measured + task success. Effect assurance remains owned by the existing invocation. + """ + + expected_effects: StateDelta = field(default_factory=StateDelta) + effect_candidates: StateDelta = field(default_factory=StateDelta) + effect_verification: EffectVerificationRequirement | None = None + scene_dependency_monitor_until: Mapping[str, int] = field(default_factory=dict) + scene_dependency_end_segment: str | None = None + + def __post_init__(self) -> None: + for name in ("expected_effects", "effect_candidates"): + value = getattr(self, name) + if type(value) is not StateDelta: + raise TypeError(f"{name} must be a StateDelta") + object.__setattr__(self, name, value.snapshot()) + if self.effect_verification is not None: + if type(self.effect_verification) is not EffectVerificationRequirement: + raise TypeError( + "effect_verification must be EffectVerificationRequirement" + ) + object.__setattr__( + self, "effect_verification", self.effect_verification.snapshot() + ) + bounds = dict(self.scene_dependency_monitor_until) + for entity, bound in bounds.items(): + _name(entity) + if type(bound) is not int or bound < 0: + raise ValueError( + "Scene monitor bounds must be nonnegative integer indices" + ) + if self.scene_dependency_end_segment is not None: + _name(self.scene_dependency_end_segment) + object.__setattr__( + self, "scene_dependency_monitor_until", MappingProxyType(bounds) + ) diff --git a/embodichain/lab/sim/motion/expansion/__init__.py b/embodichain/lab/sim/motion/expansion/__init__.py index 1e8c682a3..ea803db8d 100644 --- a/embodichain/lab/sim/motion/expansion/__init__.py +++ b/embodichain/lab/sim/motion/expansion/__init__.py @@ -95,6 +95,7 @@ validate_motion_limits, via_points, ) +from .real2sim import Real2SimSourceAdapter from .session import ExpansionSession from .profile import load_expansion_profile from .source import ( @@ -105,6 +106,7 @@ ) __all__ = [ + "Real2SimSourceAdapter", "SPATIAL_METHODS", "CandidateIdentity", "CandidateRecipe", diff --git a/embodichain/lab/sim/motion/expansion/real2sim.py b/embodichain/lab/sim/motion/expansion/real2sim.py new file mode 100644 index 000000000..51848a9ad --- /dev/null +++ b/embodichain/lab/sim/motion/expansion/real2sim.py @@ -0,0 +1,129 @@ +# ---------------------------------------------------------------------------- +# 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 qpos sources adapted to the existing source-neutral candidate contracts.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import ClassVar +import math + +import torch + +from embodichain.data_pipeline.real2sim.normalization import NormalizedReal2SimUnit +from .contracts import TrajectoryPhase, TrajectoryTemplate +from .source import SourceContext + +__all__ = ["Real2SimSourceAdapter"] + + +@dataclass(frozen=True) +class Real2SimSourceAdapter: + """Export a reviewed normalized unit without owning execution or persistence. + + Editable free windows are inset by one sample at both ends. This leaves the + source phase endpoints and adjacent central-difference stencils untouched. + Full semantic phases remain in the normalized unit, not these narrower + augmentation windows. V1 permits only same-grid joint_residual expansion. + + Unverified inputs may be normalized for offline diagnostics, but cannot + enter this generation adapter. Physical success still requires host checks. + """ + + kind: ClassVar[str] = "real2sim_qpos" + + def export_template( + self, + source: NormalizedReal2SimUnit, + *, + context: SourceContext, + ) -> TrajectoryTemplate: + """Export an owned full-joint template with conservative permissions. + + Args: + source: Audited qpos-authoritative unit on the host control clock. + context: Coordinator-owned source identity and restored scene case. + + Returns: + A template ready for the existing CandidateCoordinator. This method + does not assign candidate IDs, restore scenes or write datasets. + """ + if not isinstance(source, NormalizedReal2SimUnit) or not isinstance( + context, SourceContext + ): + raise TypeError( + "Real2SimSourceAdapter requires normalized source and SourceContext values" + ) + if source.report["authority"] != "qpos": + raise ValueError("Real2SimSourceAdapter requires qpos authority") + expected = (source.source.source_id, source.revision, source.source.unit_id) + actual = (context.source_id, context.source_revision, context.unit_id) + if actual != expected: + raise ValueError("SourceContext identity must match the normalized unit") + if not source.eligible: + raise ValueError( + "Unverified real2sim input cannot enter the qualified generation path" + ) + if not math.isclose( + context.control_dt, + float(source.report["control_dt"]), + rel_tol=0.0, + abs_tol=1e-12, + ): + raise ValueError( + "Source and host control clocks differ; normalize again explicitly" + ) + # Existing joint operators do not project arbitrary coupling graphs. + # Keep coupled drivers locked rather than producing inconsistent full states. + coupled_parents = {m["parent"] for m in source.report["mimics"]} + expanded = {source.joint_names[i] for i in source.expandable_joint_indices} + if expanded & coupled_parents: + raise ValueError("V1 expansion must keep mimic-driver groups locked") + phases = [] + for phase in source.phases: + start, stop = phase.start, phase.stop + operators = phase.allowed_operators + if operators: + if set(operators) != {"joint_residual"}: + raise ValueError( + "V1 real2sim permits only same-grid joint_residual expansion" + ) + if not source.expandable_joint_indices: + raise ValueError( + "Expansion requires explicitly authorized independent joints" + ) + start += 1 + stop -= 1 + if stop - start < 3: + raise ValueError( + "Free phase is too short for derivative-protected expansion" + ) + phases.append( + TrajectoryPhase(phase.name, start, stop, phase.kind, operators) + ) + allowed = tuple(sorted({op for p in phases for op in p.allowed_operators})) + return TrajectoryTemplate( + source_id=context.source_id, + source_revision=context.source_revision, + template_id=context.unit_id, + joint_names=source.joint_names, + positions=torch.from_numpy(source.positions.copy()), + dt=torch.from_numpy(source.dt.copy()), + phases=tuple(phases), + allowed_operators=allowed, + controlled_joint_indices=source.expandable_joint_indices, + ) diff --git a/embodichain/lab/task_program/integrations/__init__.py b/embodichain/lab/task_program/integrations/__init__.py index 3edaf84ee..e780cecfc 100644 --- a/embodichain/lab/task_program/integrations/__init__.py +++ b/embodichain/lab/task_program/integrations/__init__.py @@ -77,7 +77,21 @@ default_simulation_settle_presets, ) +from .expert_sources import ( + ExpertSourceProvider, + ExpertSourceProviderFactory, + ExpertSourceSelection, + ExpertSourceRegistry, +) +from .real2sim import Real2SimExpertSourceProvider, Real2SimExpertSourceProviderFactory + __all__ = [ + "ExpertSourceProvider", + "ExpertSourceProviderFactory", + "ExpertSourceSelection", + "ExpertSourceRegistry", + "Real2SimExpertSourceProvider", + "Real2SimExpertSourceProviderFactory", "AntipodalGraspAffordanceBinding", "ContainerAffordanceBinding", "ControlPartCommandPreset", diff --git a/embodichain/lab/task_program/integrations/_configured_composition.py b/embodichain/lab/task_program/integrations/_configured_composition.py index 834704191..74e6d0691 100644 --- a/embodichain/lab/task_program/integrations/_configured_composition.py +++ b/embodichain/lab/task_program/integrations/_configured_composition.py @@ -45,7 +45,7 @@ } ) _TASK_RUNTIME_SERVICE_FIELDS = _RUNTIME_SERVICE_FIELDS | frozenset( - {"grasp_pose_generator_overrides"} + {"grasp_pose_generator_overrides", "expert_source_providers"} ) _GRASP_GENERATOR_OVERRIDE_FIELDS = frozenset( { @@ -236,6 +236,11 @@ def _merge_runtime_services( evidence = task_evidence if task_evidence is not None else skill_profile_evidence if evidence is not None: result["control_part_evidence"] = deepcopy(evidence) + if "expert_source_providers" in task: + result["expert_source_providers"] = _owned_sequence( + task["expert_source_providers"], + path="integration.runtime_services.expert_source_providers", + ) return result @@ -282,6 +287,22 @@ def _resolve_task_program_components( ), optional=frozenset({"runtime_services"}), ) + # Data references belong to the integration file, independently of process cwd. + integration = deepcopy(dict(integration)) + services = integration.get("runtime_services", {}) + if isinstance(services, Mapping) and "expert_source_providers" in services: + services = deepcopy(dict(services)) + integration["runtime_services"] = services + for source in services["expert_source_providers"]: + if isinstance(source, dict) and "bundle_path" in source: + source["bundle_path"] = str( + _component_path( + source["bundle_path"], + base_dir=integration_path.parent, + path="expert source bundle", + suffixes=frozenset({".json"}), + ) + ) policy = _mapping( _load_yaml_component(policy_path, field_name="execution policy"), path="execution policy", diff --git a/embodichain/lab/task_program/integrations/catalog.py b/embodichain/lab/task_program/integrations/catalog.py index a2a064098..847851d5b 100644 --- a/embodichain/lab/task_program/integrations/catalog.py +++ b/embodichain/lab/task_program/integrations/catalog.py @@ -106,6 +106,8 @@ SceneReferenceRole, ) from embodichain.lab.gym.envs.task_program.bridge import RuntimeTransportActionEncoder +from .expert_sources import ExpertSourceRegistry + from .extensions import ( ControlPartEvidenceProviderDeclaration, ControlPartEvidenceProviderFactory, @@ -122,7 +124,7 @@ from .simulation import SimulationRobotSkillProfileBinding, SimulationSceneBinding from .simulation.policies import default_simulation_settle_presets -_CATALOG_FINGERPRINT_SCHEMA_VERSION = 4 +_CATALOG_FINGERPRINT_SCHEMA_VERSION = 6 _POST_POLICY_KINDS = frozenset({"wait_stable"}) _VALIDATOR_KINDS = frozenset( { @@ -480,6 +482,8 @@ class TaskProgramIntegrationCatalog: compare=False, ) + expert_sources: ExpertSourceRegistry = field(default_factory=ExpertSourceRegistry) + def __post_init__(self) -> None: for field_name in ("scene_registry_id", "robot_profile_id"): value = getattr(self, field_name) @@ -584,6 +588,10 @@ def validate_integration( f"presets are {sorted(self.robot_profile.presets)}." ) + self.expert_sources.validate_policy( + self.robot_profile.presets[integration.runtime_preset] + ) + def validate_semantic_call( self, call: SemanticCallCfg, @@ -731,6 +739,7 @@ def preflight(self, program: TaskProgramCfg) -> CompiledTaskProgram: affordance=destination, path=call.source_path, ) + self.expert_sources.validate_program(compiled) return compiled def validate_engine(self, engine: AtomicActionEngine) -> None: @@ -952,10 +961,23 @@ def _registration_payload( parallel_safety_factory: ParallelCommandSafetyValidatorFactory | None, control_part_evidence_factory: ControlPartEvidenceProviderFactory | None, registered_semantic_lowerer_factories: tuple[RegisteredSemanticLowererFactory, ...], + expert_sources: ExpertSourceRegistry, ) -> dict[str, object]: """Build the versioned canonical fingerprint payload.""" return { "schema_version": _CATALOG_FINGERPRINT_SCHEMA_VERSION, + "expert_sources": { + "selections": expert_sources.selections, + "providers": tuple( + { + "provider_id": f.provider_id, + "revision": f.revision, + "capabilities": f.capabilities, + "factory": _provider_fingerprint_declaration(f), + } + for f in expert_sources.factories + ), + }, "scene_binding": scene_binding, "scene_manifest": scene.entries, "robot_profile_binding": robot_profile_binding, @@ -1085,6 +1107,7 @@ class SimulationTaskProgramRegistration: registered_semantic_lowerer_factories: tuple[ RegisteredSemanticLowererFactory, ... ] = () + expert_sources: ExpertSourceRegistry = field(default_factory=ExpertSourceRegistry) catalog: TaskProgramIntegrationCatalog = field(init=False) _parallel_safety_validator_history: list[ParallelCommandSafetyValidator] = field( init=False, @@ -1118,6 +1141,10 @@ class SimulationTaskProgramRegistration: ) def __post_init__(self) -> None: + if type(self.expert_sources) is not ExpertSourceRegistry: + raise TypeError("expert_sources must be exactly ExpertSourceRegistry") + for selected in self.expert_sources.selections: + self.call_catalog.discover(selected.call_id) if type(self.scene_binding) is not SimulationSceneBinding: raise TypeError("scene_binding must be exactly SimulationSceneBinding.") if type(self.robot_profile_binding) is not SimulationRobotSkillProfileBinding: @@ -1153,6 +1180,12 @@ def __post_init__(self) -> None: scene = self.scene_binding.declare() profile = self.robot_profile_binding.declare() + for factory in self.expert_sources.factories: + expected_profile = getattr(factory, "target_profile_id", profile.profile_id) + if expected_profile != profile.profile_id: + raise ValueError( + "Expert source target profile does not match the registered embodiment" + ) _validate_standard_effect_monitors(profile) _validate_standard_tracking_metrics(profile) extensions = build_standard_extension_declarations( @@ -1202,6 +1235,7 @@ def __post_init__(self) -> None: fingerprint = _digest( _registration_payload( + expert_sources=self.expert_sources, scene_binding=self.scene_binding, scene=scene, robot_profile_binding=self.robot_profile_binding, @@ -1239,6 +1273,7 @@ def __post_init__(self) -> None: registered_semantic_lowerer_declarations=( extensions.registered_semantic_lowerers ), + expert_sources=self.expert_sources, fingerprint=fingerprint, _required_skills=required_skills, ), @@ -1285,6 +1320,7 @@ def assert_unchanged(self) -> None: ) current = _digest( _registration_payload( + expert_sources=self.expert_sources, scene_binding=self.scene_binding, scene=scene, robot_profile_binding=self.robot_profile_binding, diff --git a/embodichain/lab/task_program/integrations/configured.py b/embodichain/lab/task_program/integrations/configured.py index badc942ae..cb82c3e96 100644 --- a/embodichain/lab/task_program/integrations/configured.py +++ b/embodichain/lab/task_program/integrations/configured.py @@ -35,6 +35,7 @@ _PourLowererFactory, _PushObjectLowererFactory, ) +from .expert_sources import ExpertSourceRegistry, ExpertSourceSelection from .catalog import ( SimulationTaskProgramRegistration, ) @@ -2088,6 +2089,7 @@ class _DecodedRuntimeServices: ..., ] = () control_part_evidence: _JointPositionConstraintEvidenceProviderFactory | None = None + expert_sources: ExpertSourceRegistry = ExpertSourceRegistry() def _decode_runtime_services(value: object) -> _DecodedRuntimeServices: @@ -2103,6 +2105,7 @@ def _decode_runtime_services(value: object) -> _DecodedRuntimeServices: "handover_pose_providers", "registered_semantic_lowerers", "control_part_evidence", + "expert_source_providers", } ), ) @@ -2158,7 +2161,78 @@ def _decode_runtime_services(value: object) -> _DecodedRuntimeServices: path=f"{path}.control_part_evidence", ) ) + source_factories = [] + source_selections = [] + for index, item in enumerate( + _sequence( + config.get("expert_source_providers", ()), + path=f"{path}.expert_source_providers", + ) + ): + source_path = f"{path}.expert_source_providers[{index}]" + selected = _mapping( + item, + path=source_path, + required=frozenset( + { + "call_id", + "provider_id", + "revision", + "call_index", + "bundle_path", + "bundle_sha256", + "target_profile_id", + "target_asset_sha256", + } + ), + optional=frozenset({"required_capabilities"}), + ) + if selected["call_id"] != "pick": + raise ValueError(f"{source_path}: real2sim_qpos v1 supports only pick") + if (selected["provider_id"], selected["revision"]) != ("real2sim_qpos", "1"): + raise ValueError( + f"{source_path}: unregistered configured expert source version" + ) + from .real2sim import Real2SimExpertSourceProviderFactory + + source_factories.append( + Real2SimExpertSourceProviderFactory( + bundle_path=_identifier( + selected["bundle_path"], path=f"{source_path}.bundle_path" + ), + bundle_sha256=_identifier( + selected["bundle_sha256"], path=f"{source_path}.bundle_sha256" + ), + target_profile_id=_identifier( + selected["target_profile_id"], + path=f"{source_path}.target_profile_id", + ), + target_asset_sha256=_identifier( + selected["target_asset_sha256"], + path=f"{source_path}.target_asset_sha256", + ), + ) + ) + source_selections.append( + ExpertSourceSelection( + call_id=selected["call_id"], + provider_id=selected["provider_id"], + revision=selected["revision"], + call_index=selected["call_index"], + required_capabilities=tuple( + _sequence( + selected.get( + "required_capabilities", ("replay", "phase_annotation") + ), + path=f"{source_path}.required_capabilities", + ) + ), + ) + ) return _DecodedRuntimeServices( + expert_sources=ExpertSourceRegistry( + tuple(source_factories), tuple(source_selections) + ), grasp_pose_generators=generators, handover_pose_providers=handover, registered_semantic_lowerers=lowerers, @@ -2235,6 +2309,7 @@ def _decode_configured_task_program_integration( handover_pose_providers=services.handover_pose_providers, control_part_evidence_factory=services.control_part_evidence, registered_semantic_lowerer_factories=(services.registered_semantic_lowerers), + expert_sources=services.expert_sources, ) grasp_fingerprint = tuple( { diff --git a/embodichain/lab/task_program/integrations/environment.py b/embodichain/lab/task_program/integrations/environment.py index af34ae33f..c8ff94a8c 100644 --- a/embodichain/lab/task_program/integrations/environment.py +++ b/embodichain/lab/task_program/integrations/environment.py @@ -75,6 +75,7 @@ from embodichain.lab.task_program.runtime.executor import SemanticCallExecutor from embodichain.lab.task_program.runtime.expansion import ( TaskProgramPlanTransformFactory, + TaskProgramNominalPlanFactory, ) from embodichain.lab.task_program.semantics.scene import SceneRegistry @@ -300,6 +301,7 @@ class TaskProgramRuntimeAssembly: parallel_safety_validator: ParallelCommandSafetyValidator | None runtime: SemanticCallExecutor plan_transform_factory: TaskProgramPlanTransformFactory | None = None + nominal_plan_factory: TaskProgramNominalPlanFactory | None = None expansion_trace_provider: CombinedTaskProgramExpansionFactory | None = None @@ -572,17 +574,22 @@ def compile(self, program: TaskProgramCfg) -> CompiledTaskProgram: def assemble_runtime( self, integration: TaskProgramIntegrationCfg, + *, + nominal_plan_factory: TaskProgramNominalPlanFactory | None = None, ) -> TaskProgramRuntimeAssembly: """Create a fresh fully connected semantic runtime. Args: integration: Exact scene, profile, and runtime-preset selection. + nominal_plan_factory: Optional trusted, call-scoped nominal source factory. Returns: Owned assembly containing every validated runtime boundary. """ semantic = self._assemble_semantic_components(integration) - return self._assemble_execution_runtime(semantic) + return self._assemble_execution_runtime( + semantic, nominal_plan_factory=nominal_plan_factory + ) def _assemble_semantic_components( self, @@ -670,6 +677,7 @@ def _assemble_execution_runtime( semantic: SemanticExecutorComponents, *, plan_transform_factory: TaskProgramPlanTransformFactory | None = None, + nominal_plan_factory: TaskProgramNominalPlanFactory | None = None, expansion_trace_provider: CombinedTaskProgramExpansionFactory | None = None, ) -> TaskProgramRuntimeAssembly: """Attach live observation, evidence, command, and runtime boundaries.""" @@ -683,6 +691,12 @@ def _assemble_execution_runtime( "plan_transform_factory must implement " "TaskProgramPlanTransformFactory or be None." ) + if nominal_plan_factory is not None and not isinstance( + nominal_plan_factory, TaskProgramNominalPlanFactory + ): + raise TypeError( + "nominal_plan_factory must implement TaskProgramNominalPlanFactory or be None." + ) if expansion_trace_provider is not None and not isinstance( expansion_trace_provider, CombinedTaskProgramExpansionFactory, @@ -705,6 +719,26 @@ def _assemble_execution_runtime( "create_planning_observation_provider() must return a port " "implementing both ObservationProvider and CurrentQposProvider." ) + if ( + self._registration is not None + and self._registration.expert_sources.selections + ): + if nominal_plan_factory is not None: + raise ValueError( + "Registered expert sources cannot be overridden by a Python nominal factory" + ) + from embodichain.lab.sim.atomic_actions import SceneProvider + + scene_provider = getattr(observation_provider, "scene_provider", None) + if not isinstance(scene_provider, SceneProvider): + raise TypeError( + "Registered expert sources require the existing planning scene provider" + ) + nominal_plan_factory = self._registration.expert_sources.create_runtime( + scene_provider=scene_provider, + robot_profile=semantic.robot_profile, + control_dt=self._step_dt, + ) try: selected_preset = semantic.robot_profile.presets[ semantic.integration.runtime_preset @@ -775,6 +809,7 @@ def _assemble_execution_runtime( clock=clock, runner_cfg=deepcopy(selected_runner_cfg), plan_transform_factory=plan_transform_factory, + nominal_plan_factory=nominal_plan_factory, ) parallel_safety_validator = self._parallel_safety_validator if ( @@ -822,6 +857,7 @@ def _assemble_execution_runtime( parallel_safety_validator=parallel_safety_validator, runtime=runtime, plan_transform_factory=plan_transform_factory, + nominal_plan_factory=nominal_plan_factory, expansion_trace_provider=expansion_trace_provider, ) @@ -830,6 +866,7 @@ def create_bridge( program: CompiledTaskProgram, *, expansion_profile: CombinedExpansionProfile | None = None, + nominal_plan_factory: TaskProgramNominalPlanFactory | None = None, candidate_index: int = 0, recipe_indices: tuple[int, ...] | None = None, ) -> TaskProgramDemoBridge: @@ -837,12 +874,28 @@ def create_bridge( Args: program: Program compiled for this adapter's exact integration IDs. + nominal_plan_factory: Optional trusted nominal source for sequential calls. Returns: Lazy bridge sharing one newly assembled runtime, clock, and sink. """ if type(program) is not CompiledTaskProgram: raise TypeError("program must be exactly CompiledTaskProgram.") + if nominal_plan_factory is not None and any( + segment.parallel_block is not None for segment in program.iter_segments() + ): + raise ValueError( + "External nominal sources currently require sequential Task Program segments." + ) + if ( + self._registration is not None + and self._registration.expert_sources.selections + ): + self._registration.expert_sources.validate_program(program) + if expansion_profile is not None: + raise ValueError( + "Registered source expansion must use admitted phase permissions; configured combined expansion is not yet supported" + ) self._validate_selection(program.integration) self._preflight_program_surfaces(program) semantic = self._assemble_semantic_components(program.integration) @@ -875,6 +928,7 @@ def create_bridge( assembly = self._assemble_execution_runtime( semantic, plan_transform_factory=expansion, + nominal_plan_factory=nominal_plan_factory, expansion_trace_provider=expansion, ) return TaskProgramDemoBridge( diff --git a/embodichain/lab/task_program/integrations/expert_sources.py b/embodichain/lab/task_program/integrations/expert_sources.py new file mode 100644 index 000000000..c8ed7d91a --- /dev/null +++ b/embodichain/lab/task_program/integrations/expert_sources.py @@ -0,0 +1,312 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Registered nominal expert sources shared by Task Program integrations. + +Factories provide plans; Task Program retains execution, effects, recovery and +persistence. Declarations contain stable IDs and immutable parameters only. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from weakref import WeakValueDictionary +from typing import ClassVar, Protocol, runtime_checkable, TYPE_CHECKING + +from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + ActionPlan, + PlanningContext, + ResolvedActionRequest, + SceneProvider, + NominalPlanProvider, +) +from embodichain.lab.task_program.runtime.expansion import ( + TaskProgramPlanRequest, + TaskProgramNominalPlanFactory, +) +from embodichain.lab.task_program.semantics.profiles import ( + RobotSkillProfile, + SkillPolicyPreset, +) +from .extensions import validate_immutable_extension_declaration + +if TYPE_CHECKING: + from embodichain.lab.task_program.compiler.program import CompiledTaskProgram + +__all__ = [ + "ExpertSourceProvider", + "ExpertSourceProviderFactory", + "ExpertSourceSelection", + "ExpertSourceRegistry", +] + + +def _id(value: str) -> None: + if type(value) is not str or not value or value.strip() != value: + raise ValueError("Expert source IDs must be nonempty stripped strings") + + +@runtime_checkable +class ExpertSourceProvider(Protocol): + """Weak-referenceable call-local source of a side-effect-free nominal plan. + + The runtime tracks provider identity without owning completed recordings. + Providers using slots must include weak-reference support. + """ + + def plan( + self, request: ResolvedActionRequest, context: PlanningContext + ) -> ActionPlan: + """Return a complete plan or raise; source failure never falls back.""" + + +@runtime_checkable +class ExpertSourceProviderFactory(Protocol): + """Immutable, catalog-owned declaration of a fresh source provider.""" + + provider_id: ClassVar[str] + revision: ClassVar[str] + capabilities: ClassVar[frozenset[str]] + + def create( + self, + *, + call_id: str, + scene_provider: SceneProvider, + robot_profile: RobotSkillProfile, + control_dt: float, + engine: AtomicActionEngine, + ) -> ExpertSourceProvider: + """Construct a fresh provider for the selected call and host clock.""" + + +@dataclass(frozen=True) +class ExpertSourceSelection: + """Bind one expanded call occurrence to an exact registered source version. + + Args: + call_id: Semantic Call kind, for example pick. + provider_id: Stable registered provider ID. + revision: Exact provider implementation revision. + call_index: Global expanded program call index; never a kind-wide wildcard. + required_capabilities: Explicit generation capabilities required by this binding. + """ + + call_id: str + provider_id: str + revision: str + call_index: int = 0 + required_capabilities: tuple[str, ...] = ("replay", "phase_annotation") + + def __post_init__(self) -> None: + for value in (self.call_id, self.provider_id, self.revision): + _id(value) + if type(self.call_index) is not int or self.call_index < 0: + raise ValueError("Source call_index must be a nonnegative integer") + for value in self.required_capabilities: + _id(value) + if type(self.required_capabilities) is not tuple or len( + set(self.required_capabilities) + ) != len(self.required_capabilities): + raise ValueError("Source capabilities must be a unique immutable tuple") + + +@dataclass(frozen=True) +class ExpertSourceRegistry: + """Exact immutable factories and source selections in one task catalog. + + Args: + factories: Immutable factory instances, keyed by type-declared ID/revision. + selections: Exact semantic call occurrences using these factories. + """ + + factories: tuple[ExpertSourceProviderFactory, ...] = () + selections: tuple[ExpertSourceSelection, ...] = () + + def __post_init__(self) -> None: + if type(self.factories) is not tuple or type(self.selections) is not tuple: + raise TypeError("Expert source registration requires immutable tuples") + keys = [] + for factory in self.factories: + if not isinstance(factory, ExpertSourceProviderFactory): + raise TypeError("Expert source factory does not implement the protocol") + validate_immutable_extension_declaration( + factory, field_name="expert source factory" + ) + for name in ("provider_id", "revision"): + _id(getattr(type(factory), name, None)) + caps = getattr(type(factory), "capabilities", None) + if type(caps) is not frozenset: + raise TypeError( + "Factory capabilities must be a type-declared frozenset" + ) + for value in caps: + _id(value) + keys.append((factory.provider_id, factory.revision)) + if len(set(keys)) != len(keys): + raise ValueError("Duplicate expert source provider version") + indices = [] + for selected in self.selections: + if type(selected) is not ExpertSourceSelection: + raise TypeError("Invalid expert source selection") + key = selected.provider_id, selected.revision + if key not in keys: + raise ValueError(f"Unregistered expert source {key}") + factory = self.factories[keys.index(key)] + missing = set(selected.required_capabilities) - factory.capabilities + if missing: + raise ValueError( + f"Expert source lacks required capabilities: {sorted(missing)}" + ) + indices.append(selected.call_index) + if len(set(indices)) != len(indices): + raise ValueError("Duplicate source call occurrence") + if set(keys) != {(s.provider_id, s.revision) for s in self.selections}: + raise ValueError("Every registered expert source must be selected") + + def validate_policy(self, preset: SkillPolicyPreset) -> None: + """Reject replay-incompatible recovery policies before simulator construction. + + Args: + preset: Selected Task Program execution preset. + """ + for selected in self.selections: + factory = self._factory(selected) + recovery = preset.recovery_policy + if "replanning" not in factory.capabilities and ( + recovery.max_replans or recovery.max_action_retries + ): + raise ValueError( + "Replay-only expert sources require max_replans=0 and max_action_retries=0" + ) + workflow = getattr(preset, "workflow_recovery_policy", None) + if ( + "replanning" not in factory.capabilities + and workflow is not None + and workflow.max_recovery_attempts + ): + raise ValueError( + "Replay-only expert sources require workflow recovery to be disabled" + ) + if ( + "replanning" not in factory.capabilities + and preset.motion_policy.dynamic_collision_mode.value == "required" + ): + raise ValueError( + "Replay-only sources cannot satisfy dynamic collision replanning" + ) + + def validate_program(self, program: CompiledTaskProgram) -> None: + """Validate exact call coverage and supported sequencing before assembly. + + Args: + program: Immutable compiled Task Program to execute. + """ + occurrences = { + call.call_index: call.call.semantic_id + for segment in program.iter_segments() + for call in segment.calls + } + for selected in self.selections: + if occurrences.get(selected.call_index) != selected.call_id: + raise ValueError( + "Expert source selection does not match an expanded program call" + ) + if self.selections and any( + segment.parallel_block is not None for segment in program.iter_segments() + ): + raise ValueError( + "External expert sources currently require sequential program segments" + ) + + def _factory(self, selected: ExpertSourceSelection) -> ExpertSourceProviderFactory: + return next( + f + for f in self.factories + if (f.provider_id, f.revision) == (selected.provider_id, selected.revision) + ) + + def create_runtime( + self, + *, + scene_provider: SceneProvider, + robot_profile: RobotSkillProfile, + control_dt: float, + ) -> TaskProgramNominalPlanFactory: + """Create a fresh source selector for one assembled Task Program runtime. + + Args: + scene_provider: Existing scene snapshot boundary. + robot_profile: Existing registered embodiment profile. + control_dt: Host-owned command period. + + Returns: + Call-scoped nominal-plan factory using this exact registration. + """ + return _RegisteredNominalSources( + self, scene_provider, robot_profile, control_dt + ) + + +class _RegisteredNominalSources: + def __init__(self, registry, scene_provider, robot_profile, control_dt): + self.registry, self.scene_provider, self.robot_profile, self.control_dt = ( + registry, + scene_provider, + robot_profile, + control_dt, + ) + self._providers: WeakValueDictionary[int, ExpertSourceProvider] = ( + WeakValueDictionary() + ) + + def create_nominal_plan_provider( + self, request: TaskProgramPlanRequest, *, engine: AtomicActionEngine + ) -> NominalPlanProvider | None: + selected = next( + ( + s + for s in self.registry.selections + if s.call_index == request.analysis_call_index + ), + None, + ) + if selected is None: + return None + if selected.call_id != request.call.semantic_id: + raise ValueError( + "Source call kind does not match selected program occurrence" + ) + factory = self.registry._factory(selected) + provider = factory.create( + call_id=selected.call_id, + scene_provider=self.scene_provider, + robot_profile=self.robot_profile, + control_dt=self.control_dt, + engine=engine, + ) + if not isinstance(provider, ExpertSourceProvider): + raise TypeError("Factory returned an invalid expert source provider") + if self._providers.get(id(provider)) is provider: + raise ValueError("Factory reused live expert source state") + try: + self._providers[id(provider)] = provider + except TypeError as error: + raise TypeError( + "Expert source providers must support weak references" + ) from error + return provider.plan diff --git a/embodichain/lab/task_program/integrations/real2sim.py b/embodichain/lab/task_program/integrations/real2sim.py new file mode 100644 index 000000000..544720350 --- /dev/null +++ b/embodichain/lab/task_program/integrations/real2sim.py @@ -0,0 +1,336 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Qpos-authoritative expert provider with Skill-owned plan construction.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import ClassVar +import math +import hashlib +from pathlib import Path + +import torch + +from embodichain.data_pipeline.real2sim.bundle import load_real2sim_bundle +from embodichain.lab.sim.atomic_actions.trajectory_assembly import SkillTrajectoryInput +from embodichain.data_pipeline.real2sim.normalization import NormalizedReal2SimUnit +from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + PlanningContext, + ResolvedActionRequest, + SceneProvider, +) +from embodichain.lab.sim.atomic_actions.plans import ( + ActionPlan, + PlannerDiagnostics, + TimedTrajectory, + TrajectorySegment, +) +from embodichain.lab.sim.motion.expansion import ( + Real2SimSourceAdapter, + SourceContext, + SceneCase, +) +from embodichain.lab.task_program.semantics.profiles import RobotSkillProfile + +__all__ = ["Real2SimExpertSourceProviderFactory", "Real2SimExpertSourceProvider"] + + +@dataclass(frozen=True) +class Real2SimExpertSourceProviderFactory: + """Allowlisted source declaration pinned by a complete bundle digest. + + Args: + bundle_path: Non-executable real2sim bundle declaration. + bundle_sha256: Required content digest of that declaration. + target_profile_id: Exact reviewed embodiment profile bound to this source. + target_asset_sha256: Pinned resolved URDF/USD definition of the target robot. + """ + + provider_id: ClassVar[str] = "real2sim_qpos" + revision: ClassVar[str] = "1" + capabilities: ClassVar[frozenset[str]] = frozenset( + {"replay", "phase_annotation", "effect_validation"} + ) + bundle_path: str + bundle_sha256: str + target_profile_id: str + target_asset_sha256: str + + def __post_init__(self) -> None: + if ( + not self.target_profile_id + or self.target_profile_id.strip() != self.target_profile_id + ): + raise ValueError("An exact target profile ID is required") + if len(self.target_asset_sha256) != 64 or any( + c not in "0123456789abcdef" for c in self.target_asset_sha256 + ): + raise ValueError("A lowercase target asset SHA-256 is required") + # Admission is provider-free and happens before the environment is built. + unit = load_real2sim_bundle( + self.bundle_path, expected_sha256=self.bundle_sha256 + ) + _segment_lengths(unit) + _template(unit) + + def create( + self, + *, + call_id: str, + scene_provider: SceneProvider, + robot_profile: RobotSkillProfile, + control_dt: float, + engine: AtomicActionEngine, + ) -> Real2SimExpertSourceProvider: + """Create an owned B=1 Pick source on the host clock. + + Args: + call_id: Selected semantic kind; v1 supports pick. + scene_provider: Existing scene ownership boundary. + robot_profile: Existing embodiment profile. + control_dt: Fixed environment command interval. + engine: Existing Skill construction and validation owner. + + Returns: + Fresh fixed-qpos provider with independently loaded source evidence. + """ + if call_id != "pick": + raise ValueError("real2sim_qpos v1 supports one Pick call") + if robot_profile.profile_id != self.target_profile_id: + raise ValueError( + "Source target profile differs from the selected embodiment" + ) + asset_path = Path(engine.robot.cfg.fpath) + if ( + hashlib.sha256(asset_path.read_bytes()).hexdigest() + != self.target_asset_sha256 + ): + raise ValueError( + "Source target asset digest differs from the live robot model" + ) + unit = load_real2sim_bundle( + self.bundle_path, expected_sha256=self.bundle_sha256 + ) + if not math.isclose( + control_dt, float(unit.report["control_dt"]), abs_tol=1e-12, rel_tol=0 + ): + raise ValueError("Source control clock differs from the environment") + return Real2SimExpertSourceProvider( + unit, engine, robot_profile.profile_id, self.target_asset_sha256 + ) + + +def _segment_lengths(unit) -> dict[str, int]: + # These are the full semantic phases, never the inset augmentation windows. + phases = unit.phases + if ( + tuple(p.name for p in phases) != ("approach", "close", "lift") + or phases[0].start != 0 + or phases[-1].stop != len(unit.positions) + ): + raise ValueError( + "A recorded Pick source must have complete approach/close/lift phases" + ) + if phases[1].kind != "contact" or phases[2].kind != "hold": + raise ValueError( + "Recorded Pick close/lift phases must be protected contact/hold" + ) + # Interpolation-only gaps stay locked by the source adapter. Attribute them + # to the preceding semantic segment, preserving every hard event sample. + stops = [phases[1].start, phases[2].start, len(unit.positions)] + return {p.name: stop - p.start for p, stop in zip(phases, stops)} + + +def _template(unit): + return Real2SimSourceAdapter().export_template( + unit, + context=SourceContext( + unit.source.source_id, + unit.revision, + unit.source.unit_id, + SceneCase( + "source_admission", + "declared_initial", + "source_scene", + "pick", + "source_profile", + ), + float(unit.report["control_dt"]), + ), + ) + + +@dataclass +class Real2SimExpertSourceProvider: + """One owned recording; no execution, reset or persistence lifecycle. + + Args: + unit: Verified normalized source unit. + engine: Existing engine that owns Skill plan construction. + profile_id: Selected embodiment profile identity for provenance. + asset_sha256: Verified live model definition digest. + """ + + unit: NormalizedReal2SimUnit + engine: AtomicActionEngine + profile_id: str + asset_sha256: str + + def plan( + self, request: ResolvedActionRequest, context: PlanningContext + ) -> ActionPlan: + """Build a full Pick plan from admitted full-joint source values. + + Args: + request: Exact resolved PickUp invocation. + context: Live initial robot and semantic scene state. + + Returns: + Complete Skill-assembled plan with source diagnostics. + """ + unit = self.unit + _template(unit) + if context.batch_size != 1: + raise ValueError("Recorded qpos v1 requires B=1") + if request.skill_id != "pick_up": + raise ValueError("Recorded Pick source requires the PickUp Skill") + if ( + request.revision != 0 + or request.recovery_policy.max_replans + or request.recovery_policy.max_action_retries + ): + raise ValueError( + "Fixed qpos replay cannot replan or retry without a fresh reset" + ) + if tuple(self.engine.robot.joint_names) != unit.joint_names: + raise ValueError( + "Source joint order differs from the complete simulation robot" + ) + if not math.isclose( + context.require_control_dt(), + float(unit.report["control_dt"]), + rel_tol=0, + abs_tol=1e-12, + ): + raise ValueError("Source control clock differs from the planning context") + q = torch.tensor( + unit.positions.copy(), + device=context.robot.qpos.device, + dtype=context.robot.qpos.dtype, + )[None] + if not torch.allclose(q[:, 0], context.robot.qpos, atol=1e-4, rtol=0): + raise ValueError( + "Restore the exact recording initial state before fixed replay" + ) + limits = self.engine.robot.get_qpos_limits(env_ids=context.env_ids) + if not isinstance(limits, torch.Tensor) or limits.shape != ( + 1, + len(unit.joint_names), + 2, + ): + raise ValueError("Live robot must expose full-joint position limits") + limits = limits.to(q) + if not torch.isfinite(limits).all() or (limits[..., 0] > limits[..., 1]).any(): + raise ValueError("Live joint position limits must be finite and ordered") + if ( + (q < limits[:, None, :, 0] - 1e-6) | (q > limits[:, None, :, 1] + 1e-6) + ).any(): + raise ValueError("Source violates live robot joint position limits") + velocity_limits = self.engine.robot.get_qvel_limits(env_ids=context.env_ids) + if ( + not isinstance(velocity_limits, torch.Tensor) + or velocity_limits.shape != (1, len(unit.joint_names)) + or not torch.isfinite(velocity_limits).all() + or (velocity_limits <= 0).any() + ): + raise ValueError( + "Live robot must expose positive finite joint velocity limits" + ) + velocities = torch.tensor( + unit.velocities.copy(), device=q.device, dtype=q.dtype + )[None] + if (velocities.abs() > velocity_limits.to(q)[:, None] + 1e-6).any(): + raise ValueError("Source violates live robot joint velocity limits") + dt = torch.tensor(unit.dt.copy(), device=q.device, dtype=q.dtype)[None] + # Central differences can cancel at reversals. Position commands must + # also be reachable within every individual arrival interval. + interval_speeds = q.diff(dim=1).abs() / dt[:, 1:, None] + if (interval_speeds > velocity_limits.to(q)[:, None] + 1e-6).any(): + raise ValueError( + "Source interval speeds violate live robot joint velocity limits" + ) + claimed = set() + for endpoint in request.binding.endpoints: + target = getattr(endpoint, "target", None) + claimed.update(getattr(target, "joint_ids", ())) + moving = ( + ((q[:, 1:] - q[:, :-1]).abs().amax(dim=(0, 1)) > 1e-8) + .nonzero() + .flatten() + .tolist() + ) + # Mimics are derived channels. The source adapter already locks their drivers. + mimics = {unit.joint_names.index(m["joint"]) for m in unit.report["mimics"]} + if set(moving) - claimed - mimics: + raise ValueError( + "Recording moves joints outside the call's authorized endpoints" + ) + trajectory = TimedTrajectory( + positions=q, + velocities=velocities, + accelerations=None, + dt=dt, + env_ids=context.env_ids, + ) + contract = self.engine.skills[request.skill_id].trajectory_contract + if contract is None: + raise ValueError( + "Selected Skill does not declare a trajectory assembly schema" + ) + segments = [] + offset = 0 + for name, length in _segment_lengths(unit).items(): + segments.append(TrajectorySegment(name, offset, offset + length)) + offset += length + segments = tuple(segments) + return self.engine.build_plan_from_trajectory( + request, + context, + SkillTrajectoryInput( + schema_id=contract.schema_id, + schema_revision=contract.revision, + trajectory=trajectory, + segments=segments, + landmarks=contract.stage_landmarks(segments, trajectory.waypoint_count), + diagnostics=PlannerDiagnostics( + backend="real2sim_qpos", + metadata={ + "source_id": unit.source.source_id, + "source_revision": unit.source.source_revision, + "normalized_revision": unit.revision, + "unit_id": unit.source.unit_id, + "authority": "qpos", + "profile_id": self.profile_id, + "target_asset_sha256": self.asset_sha256, + "normalization": unit.report_dict(), + "interphase_segment_policy": "preceding_segment_locked", + }, + ), + ), + ) diff --git a/embodichain/lab/task_program/runtime/__init__.py b/embodichain/lab/task_program/runtime/__init__.py index 4a51969b2..62d393cd7 100644 --- a/embodichain/lab/task_program/runtime/__init__.py +++ b/embodichain/lab/task_program/runtime/__init__.py @@ -19,7 +19,11 @@ from __future__ import annotations from .executor import SemanticCallExecutor -from .expansion import TaskProgramPlanRequest, TaskProgramPlanTransformFactory +from .expansion import ( + TaskProgramPlanRequest, + TaskProgramPlanTransformFactory, + TaskProgramNominalPlanFactory, +) from .parallel import ParallelTimingPolicy from .parallel_executor import ParallelSemanticExecutor from .results import SemanticExecutionResult, SemanticExecutionStatus @@ -30,6 +34,7 @@ "SemanticCallExecutor", "TaskProgramPlanRequest", "TaskProgramPlanTransformFactory", + "TaskProgramNominalPlanFactory", "SemanticExecutionResult", "SemanticExecutionStatus", ] diff --git a/embodichain/lab/task_program/runtime/executor.py b/embodichain/lab/task_program/runtime/executor.py index 3ad0a02a2..4553af176 100644 --- a/embodichain/lab/task_program/runtime/executor.py +++ b/embodichain/lab/task_program/runtime/executor.py @@ -99,7 +99,11 @@ _snapshot_event, _snapshot_task_state, ) -from .expansion import TaskProgramPlanRequest, TaskProgramPlanTransformFactory +from .expansion import ( + TaskProgramPlanRequest, + TaskProgramPlanTransformFactory, + TaskProgramNominalPlanFactory, +) from ..compiler.lowering import ( GroundedHeldObjectGuard, GroundedPhaseEffectGate, @@ -267,6 +271,7 @@ def __init__( clock: ExecutionClock | None = None, runner_cfg: ExecutionRunnerCfg | None = None, plan_transform_factory: TaskProgramPlanTransformFactory | None = None, + nominal_plan_factory: TaskProgramNominalPlanFactory | None = None, ) -> None: if not isinstance(compiler, SemanticCallCompiler): raise TypeError("compiler must be a SemanticCallCompiler.") @@ -290,6 +295,13 @@ def __init__( "plan_transform_factory must implement " "TaskProgramPlanTransformFactory or be None." ) + if nominal_plan_factory is not None and not isinstance( + nominal_plan_factory, TaskProgramNominalPlanFactory + ): + raise TypeError( + "nominal_plan_factory must implement TaskProgramNominalPlanFactory or be None." + ) + self._nominal_plan_factory = nominal_plan_factory integration = compiler.integration engine = integration.engine if not isinstance(engine, AtomicActionEngine): @@ -1037,7 +1049,11 @@ def _prepare_grounded_call( "Grounded semantic call preset must own an ExecutionRunnerCfg." ) plan_transform = None - if self._plan_transform_factory is not None: + nominal_plan_provider = None + if ( + self._plan_transform_factory is not None + or self._nominal_plan_factory is not None + ): if self._workflow_id is None: raise RuntimeError("A grounded call requires an active workflow ID.") plan_request = TaskProgramPlanRequest( @@ -1047,6 +1063,7 @@ def _prepare_grounded_call( call=call, invocation=invocation, ) + if self._plan_transform_factory is not None: prepare_context = getattr( self._plan_transform_factory, "prepare_planning_context", @@ -1075,11 +1092,25 @@ def _prepare_grounded_call( raise TypeError( "create_plan_transform() must return a callable or None." ) + if self._nominal_plan_factory is not None: + nominal_plan_provider = ( + self._nominal_plan_factory.create_nominal_plan_provider( + plan_request, + engine=self._engine, + ) + ) + if nominal_plan_provider is not None and not callable( + nominal_plan_provider + ): + raise TypeError( + "create_nominal_plan_provider() must return a callable or None." + ) session = self._engine.start( (invocation,), context, eligible_mask=active_mask, plan_transform=plan_transform, + nominal_plan_provider=nominal_plan_provider, ) primed = _PrimedObservationProvider(context, self._observation_provider) runner = ExecutionRunner( diff --git a/embodichain/lab/task_program/runtime/expansion.py b/embodichain/lab/task_program/runtime/expansion.py index 78acce24d..6fc466d8c 100644 --- a/embodichain/lab/task_program/runtime/expansion.py +++ b/embodichain/lab/task_program/runtime/expansion.py @@ -25,10 +25,15 @@ ActionInvocation, AtomicActionEngine, PlanTransform, + NominalPlanProvider, ) from embodichain.lab.task_program.semantics.calls import SemanticCallSpec -__all__ = ["TaskProgramPlanRequest", "TaskProgramPlanTransformFactory"] +__all__ = [ + "TaskProgramPlanRequest", + "TaskProgramPlanTransformFactory", + "TaskProgramNominalPlanFactory", +] @dataclass(frozen=True, slots=True) @@ -75,3 +80,31 @@ def create_plan_transform( engine: AtomicActionEngine, ) -> PlanTransform | None: """Return one session-scoped transform or ``None`` for no change.""" + + +@runtime_checkable +class TaskProgramNominalPlanFactory(Protocol): + """Create a source-specific complete-plan provider for one grounded call. + + Trusted integrations own registration and capability preflight. Factories + must select exact workflow/call occurrences and return fresh call-local + providers. None means this call has no external source binding; a selected + provider must raise on failure rather than falling back to ordinary planning. + Providers own neither candidate IDs, simulator stepping nor persistence. + """ + + def create_nominal_plan_provider( + self, + request: TaskProgramPlanRequest, + *, + engine: AtomicActionEngine, + ) -> NominalPlanProvider | None: + """Select the nominal source without invoking ordinary motion planning. + + Args: + request: Stable identity and grounded invocation of the current call. + engine: Owning engine, used for existing skill plan construction. + + Returns: + A call-scoped plan provider, or None for an unbound call. + """ diff --git a/examples/real2sim/README.md b/examples/real2sim/README.md new file mode 100644 index 000000000..b6e91a4fe --- /dev/null +++ b/examples/real2sim/README.md @@ -0,0 +1,348 @@ +# Source-independent real2sim foundation + +This example exercises the reusable qpos source contracts proposed in #709. +It does not import a DROID experiment, reconstruct a scene, tune a robot, or +start a second collection lifecycle. The original recording remains external +and read-only. Ordinary tests use synthetic arrays, not this optional file. + +## Implemented slice + +- Owned raw trajectory, joint-layout, phase and optional EEF evidence values. +- Explicit unit/joint mapping, initial-state checks, declared holds and mimic derivation. +- Fixed-clock resampling with exact event samples and bidirectional timing evidence. +- A `Real2SimSourceAdapter` for the existing `CandidateCoordinator`/`ExpansionSession`. +- A complete nominal-plan injection boundary in the Atomic engine, retained during + recovery and invocation revision, and a call-scoped Task Program factory hook. +- Regression tests for bypassing ordinary planning, preserving effect verification, + deterministic variants, protected derivative boundaries and malformed data. + +The shared `ExpertSourceRegistry` now fingerprints exact provider versions and +expanded call selections in the simulation catalog. The configured +`expert_source_providers` service admits a content-pinned source bundle before +simulation construction. `Real2SimExpertSourceProviderFactory` creates fresh B=1 +Pick providers, and PickUp supplies semantics to the shared trajectory assembler. +`SkillTrajectoryContract` and `SkillTrajectoryInput` are source-independent; +ordinary PickUp and MoveJoints planning use the same assembly path. Neither source +code nor callables belong in `program.yaml`. + +Physical qualification, receipt-backed array attachments, registered-source +combined expansion, time-grid changes and `eef_pose`/`hybrid` planning remain +subsequent milestones. The existing same-grid source adapter and coordinator +remain available; configured combined expansion with a registered source is +rejected rather than changing its phase permissions. + +## Contracts + +Create `Real2SimTrajectory` and `RobotJointLayout` independently of a dataset +format, then call `normalize_real2sim(source, layout, cfg)`. The host supplies +`control_dt`; normalization never changes the simulator physics clock. + +Source timestamps strictly increase. Duplicate timestamps must be coalesced by +an explicit upstream operation that preserves event provenance. Each phase is +half-open `[start, stop)`, with protected samples at `start` and `stop - 1`. +Hard anchors round forward to control ticks. The default maximum delay is one +control tick per anchor. Incompatible short event sequences raise instead of +silently dropping contact phases. Interpolation-only gaps between remapped +phases remain unannotated and therefore locked. + +The v1 optional `EefPoseSamples` stream is aligned to source joint timestamps. +An independent EEF clock must be aligned explicitly upstream and its operation +recorded in provenance. Translations in frame/TCP transforms are in metres; +pose translation and quaternion conventions are declared separately. An +`fk_derived` pose is never reported as independent measurement evidence. + +The adapter always rejects unverified input prerequisites. Offline diagnostic +normalization is possible with `allow_unverified=True`, but that output cannot +enter the adapter. Verified prerequisites still do not imply physical success. +V1 permits only same-grid `joint_residual` and requires enough free samples to +protect both endpoints and their adjacent derivative stencils. Template free +windows are inset by one sample at each end; full semantic phase ranges remain +in the normalized unit/report and must be used by future Skill plan assembly. +Expandable mimic-driver groups are rejected until an explicit coupling-aware +candidate projection is available. Ordinary gripper groups can remain locked. + +The template source revision is the normalization content digest. The host's +`SourceContext` must match source ID, unit ID, revision and control period; +candidate IDs are still allocated only by the existing coordinator/session. + +## Batched anchored resampling + +The compute layer also accepts float32/float64 `(B,N_max,D)` tensors with +independent source lengths, clocks and anchor counts. Prepare a mapping once +when multiple candidate paths share the same timing: + +```python +from embodichain.compute.trajectory import resample_with_anchors_batched + +sampled, plan = resample_with_anchors_batched( + qpos, source_times, 0.01, anchor_indices, + source_lengths=source_lengths, anchor_counts=anchor_counts, + backend="warp", +) +sampled_variant = plan.apply(qpos_variant, backend="warp") +valid_samples = plan.valid_mask +``` + +The host prepares float64 time maps using the single-path algorithm; evaluation +preserves input dtype/device and runs on the active Torch CUDA stream. Torch is +the default backend; explicit Warp failures propagate. Source/anchor padding is +ignored, output padding holds the final valid state, and valid counts/masks must +be used for evaluation and persistence. Warp qpos autograd is unsupported; use +Torch for gradient-bearing inputs. Neither API computes task acceptance. +The registered Task Program source remains B=1. + +Run `python -m scripts.benchmark.trajectory.anchored_batch` to compare NumPy, +Torch CPU/CUDA and Warp CPU/CUDA. The report separates prepared evaluation from +mapping preparation and offline input/output transfer, alongside numerical +agreement, exact anchors and memory measurements. + +## One optional real recording + +Run from a checkout with the normal EmbodiChain dependencies installed: + +```bash +python examples/real2sim/normalize_recording.py \ + --input /path/to/recording.npz \ + --output outputs/real2sim-source-sample/run01 \ + --joint-names joint_a joint_b \ + --joint-units rad rad \ + --robot-id recorded_robot \ + --robot-revision unverified-recording-model \ + --time-unit ms --control-dt 0.01 +``` + +The NPZ defaults to `joint_positions`, `timestamps` and optional +`joint_velocities`; key names are CLI arguments. `--start-index` and +`--stop-index` select one source slice. The provided channels form a diagnostic +layout, not a verified full robot. The script deliberately marks calibration +and phase annotations unverified and cannot certify a Pick or commit a dataset +episode. It refuses to write into the source directory or overwrite its reports. + +Optional EEF evidence requires both `--eef-key` and `--eef-description`. +The JSON description supplies `reference_frame`, `observed_link`, `tcp_id`, +`reference_to_base`, `link_to_tcp`, `origin`, `quaternion_order` and +`position_unit`. Supplying a TCP pose requires an identity link-to-TCP transform; +do not apply the same tool offset twice. + +The output contains `trajectory.npz` (raw and normalized arrays) and +`normalization.json` (source digest, unit mappings, phase/timing maps and explicit +qualification limitations). These are offline diagnostic artifacts, not +`EpisodeSink` commit receipts. Gripper channels must not be reinterpreted from +their numeric range. Decode a known width/angle convention upstream, or leave +it unresolved and keep the sample diagnostic-only. + +## Focused validation + +```bash +pytest tests/compute/test_trajectory_anchors.py tests/data_pipeline/real2sim \ + tests/sim/motion/expansion/test_real2sim.py +pytest tests/sim/atomic_actions/test_engine.py tests/sim/atomic_actions/test_engine_per_env.py +pytest tests/lab/task_program/test_semantic_executor.py tests/gym/envs/task_program/test_environment.py +``` + +Tests exercise production contracts. No tutorial, experiment outputs, local +recording path or reconstructed asset is a CI dependency. + +## Four-panel kinematic video + +`render_preview.py` adds a native EmbodiChain renderer preview for the optional +Franka arm recording. It shows the source camera, nominal arm qpos, and two +seeded joint-residual variants. The small plots show **flange XY paths**, not a +calibrated gripper TCP. The source NPZ must explicitly contain seven Franka +`joint_positions` columns in `fr3_joint1`–`fr3_joint7` order and `timestamps` in +milliseconds. The source video must be 640×360 with N or N−1 frames. + +```bash +PYTHONPATH="$PWD:$PWD/extensions" python examples/real2sim/render_preview.py \ + --input /path/to/gt_action_full.npz \ + --video /path/to/camera_rgb.mp4 \ + --output outputs/real2sim-preview/run01 \ + --free-stop 60 --scale 0.025 +``` + +Choose `--free-stop` only after reviewing the recording: it is the exclusive +source sample index of a conservative free-motion prefix. For the selected +sample, 60 protects the final approach and all later contact, transport and +release motion. The protected suffix uses a conservative contact classification; +it does not claim to have finely segmented those actions. The editable window +is inset to preserve boundary positions and velocity stencils. + +The preview deliberately remains outside qualified source admission: calibration +and annotations remain unverified, `Real2SimSourceAdapter` is not relaxed, and the +example applies the existing `joint_residual` operator directly to a diagnostic +`TrajectoryTemplate`. It does not claim a `CandidateCoordinator` acceptance or +an `EpisodeSink` receipt. Synthetic production tests cover the qualified adapter +separately. + +Three render arenas display the same official Franka arm and Robotiq 2F85 assets, +assembled inside the output directory. There are no reconstructed scene objects. +The gripper defaults to fixed open; the optional explicit lookup below animates +the declared scalar convention. The mount remains an uncalibrated identity transform. Robot states are written for +kinematic inspection and rendered without physics stepping. Joint-position and +velocity limits, protected qpos/qvel and runtime qpos readback are checked; +collisions, tracking under dynamics and task success are not qualified. + +The video follows the state clock, not the camera file's nominal 30 FPS. Default +video frame i maps to source state i; `--video-state-offset -1|0|1` makes that +assumption explicit. Nearest source samples and endpoint clamping handle N−1 +frames. The offset still requires independent projection-based calibration. + +Outputs include `comparison.mp4`, review stills, `trajectories.npz` (all three +paths, qvel, FK, protected mask and render/source/video mappings), `report.json` +(provenance, asset hashes, amplitude/limit metrics and limitations), and the +assembled robot asset. `--max-frames 1` runs a bounded render smoke check. Outputs +must be in a new directory outside the original recording directory. + + +## Registered B=1 Pick integration + +A task-local integration can select a verified source using only stable IDs and +content-pinned data references: + +```yaml +runtime_services: + expert_source_providers: + - call_id: pick + call_index: 0 + provider_id: real2sim_qpos + revision: "1" + target_profile_id: my_franka_profile + target_asset_sha256: + bundle_path: ./sources/pick/bundle.json + bundle_sha256: +``` + +`call_index` selects one global expanded program occurrence. The catalog rejects +missing/mismatched occurrences, unknown versions, duplicate selections and +unsupported capabilities. Data paths resolve relative to the integration file. +The normal execution-policy component must set both `max_replans: 0` and +`max_action_retries: 0` for fixed replay. External nominal sources currently +require sequential segments. Fresh providers share the existing planning scene +provider and exact embodiment profile; Python callbacks cannot override a +registered source binding. + +Create source bundles with `write_real2sim_bundle(source, layout, cfg, directory)`. +`load_real2sim_bundle` verifies the declaration/NPZ digests and closed fields, +reconstructs raw contracts, reruns normalization and rejects unreviewed inputs by +default. Large arrays stay in the NPZ sidecar. There are no executable loaders or +pickled arrays. Data digests are rechecked when creating each live source provider. + +The initial Pick source has full `approach`, `close`, `lift` semantic phases. +Close and lift use protected `contact`/`hold` kinds. Interpolation-only gaps remain +locked by the adapter; the plan assembler attributes those gaps to the preceding +semantic segment. All source/execution phase intervals remain in diagnostics. +The provider checks B=1, exact profile/model digest/joint order, initial state, control period, +endpoint ownership and live joint-position limits. The provider calls +`AtomicActionEngine.build_plan_from_trajectory`; the selected Skill supplies +semantic effects to the existing command/tracking builder. Plans retain owned +assembly inputs so same-grid candidates rebuild commands, tracking, effects, +stages and dependency boundaries together. Geometry changes require new facts +rather than carrying stale grasp baselines. +Source failures propagate without ordinary planning fallback. A fresh reset is +required to run another fixed source attempt. + +Speed admission checks both the differentiated references and every commanded +position interval, so reversals cannot cancel a speed-limit violation. Provider +instances support weak references: the selector tracks identity without keeping +completed recordings alive, and cached provider reuse is still rejected. + +The B=1 bridge regression uses reviewed **synthetic** source values and projected +semantic completion. It proves the interface and lifecycle, not physical grasp +success. Existing verified-runtime tests continue to cover measured-effect failure. + +## Original-recording calibration audit + +`prepare_pick_sample.py` consumes only one original HDF5, independently of any +real2sim experiment code or generated scene: + +```bash +PYTHONPATH="$PWD:$PWD/extensions" python examples/real2sim/prepare_pick_sample.py \ + --hdf5 /path/to/original_sample.hdf5 \ + --output outputs/real2sim-source-sample/pick_audit \ + --upstream-revision c5737e40a6b18859b5b78dbcdbf1e3b3f5e461be \ + --close-start 97 --lift-start 106 --stop 121 +``` + +The selected source has an open approach, a closing transition at samples 97–105, +and subsequent lift. These indices are example annotations, not an automatic +contact detector. The script preserves raw gripper samples, camera matrices, +source pose and pipeline calibration metadata in a separate evidence sidecar. +It writes a diagnostic source bundle, explicit jaw-width/driver lookup and audit +report. Original files remain read-only. + +The [pinned upstream DROID implementation](https://github.com/droid-dataset/droid/blob/c5737e40a6b18859b5b78dbcdbf1e3b3f5e461be/droid/franka/robot.py) +defines gripper position as normalized closure (`1 - width/max_width`). The audit +uses this explicit convention, preserves the exported canonical values separately +in the original NPZ, and does not interpret the scalar as driver radians. The +width-to-driver table comes from the bundled Robotiq contact-pad mesh and includes +estimated endpoint extensions; it remains unverified hardware calibration. + +To animate that estimate in the diagnostic renderer, add +`--gripper-calibration /path/to/gripper_calibration.json` to the four-panel command. +This replaces the fixed-open visual policy without granting formal admission. + +The original sample still lacks independently reviewed real mount/TCP, hardware +jaw calibration, camera reprojection/frame offset and live target-joint order. +Its audit therefore reports `rejected_until_reviewed`. The near-zero source-pose +versus flange-FK residual is a conversion consistency check, not an independent +measurement-accuracy result. Do not promote these input states automatically. + +## Physical interaction and expansion preview + +`physics_preview.py` drives the same recorded arm path and two seeded residual +variants through the shared `play_joint_trajectory` host, with 1 ms Default +physics and 10 ms PD command periods. Each candidate runs sequentially at B=1 +from the same reset state. The block is dynamic; table and container proxies are +static. Object poses are written only at reset. Gravity, contact and friction +produce grasping, transport, release and settling; no object attachment or +object-trajectory playback is used. + +```bash +PYTHONPATH="$PWD:$PWD/extensions" python examples/real2sim/physics_preview.py \ + --input /path/to/gt_action_full.npz \ + --gripper-calibration /path/to/gripper_calibration.json \ + --source-video /path/to/camera_rgb.mp4 \ + --output outputs/real2sim-physics/run01 --repeats 3 +``` + +Generate the explicit gripper lookup with `prepare_pick_sample.py` first. This +example uses the selected recording's source indices 105 and 160 to anchor the +estimated grasp and release geometry. The block width follows the declared jaw +width at closure; length/height, mass, friction, container and table are shared +model estimates. The robot base is translated to 0.75 m in the preview world. +The actual measured gripper position is converted into a controller target with +an explicit 0.015 rad squeeze preload while closing. This preload is shared by +all candidates, stored separately from raw driver values, and configurable with +`--grip-preload`. It is a control assumption, not recovered real command evidence. + +Use `--no-render --cases 0 1 2 --repeats 3` for a bounded physics-only run, or select +one case while investigating failure. Parameters such as `--mass`, `--friction`, +`--grip-effort` and `--grip-stiffness` are explicit and recorded. They must remain +fixed across a comparison; changing them produces another diagnostic condition. + +The contact sensor resolves exact left/right inner-finger actor identities. The +measured proxy-scene criteria require the block to rise at least 4 cm with +continuous bilateral contact for at least 1 s, finish wholly inside the container +(using oriented box corners), and remain stable in position and orientation for +the final 1 s. During that final second the measured gripper driver must also +remain within `--release-open-tolerance` (default 0.03 rad) of its calibrated open +position, with no bilateral grasp contact. An opening command alone is insufficient. +Both hold and release durations use elapsed timestamps. Stability also limits +linear and angular speed. A motion-limit pass +or a robot reaching its endpoint does not establish this physical result. + +Outputs include per-candidate native videos, the four-panel +`physical_comparison.mp4` when `--source-video` is supplied, controller targets, +measured object/joint/flange traces, contact/impulse evidence, actor identities, +settings and per-trial results. The first repeat of each candidate is rendered; +every repeat is physically executed and evaluated. Interpolation-only and contact +motion remains protected exactly as in the kinematic diagnostic. + +These results qualify the **estimated proxy-scene diagnostic only**. The real +mount/TCP, hardware jaw calibration, camera/frame alignment and reconstructed +asset geometry remain unverified. Physical-preview success does not change the +source bundle's admission state, allocate accepted generation candidates or +produce an EpisodeSink commit receipt. The registered real2sim source still +rejects the unreviewed original recording. + +Both preview CLIs accept the shared `--viser` visualization options. diff --git a/examples/real2sim/normalize_recording.py b/examples/real2sim/normalize_recording.py new file mode 100644 index 000000000..26b05a1f4 --- /dev/null +++ b/examples/real2sim/normalize_recording.py @@ -0,0 +1,186 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Optional real-recording example; no dependency on any reconstruction experiment.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np + +from embodichain.data_pipeline.real2sim import ( + EefPoseSamples, + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, +) +from embodichain.data_pipeline.real2sim.normalization import ( + Real2SimNormalizationCfg, + normalize_real2sim, +) + +__all__: list[str] = [] + + +def main() -> None: + """Normalize one external NPZ recording as explicitly unverified evidence.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--joint-names", nargs="+", required=True) + parser.add_argument( + "--joint-units", nargs="+", choices=["rad", "deg", "m", "mm"], required=True + ) + parser.add_argument("--robot-id", required=True) + parser.add_argument("--robot-revision", required=True) + parser.add_argument("--time-unit", choices=["s", "ms", "us", "ns"], required=True) + parser.add_argument("--control-dt", type=float, required=True) + parser.add_argument("--positions-key", default="joint_positions") + parser.add_argument("--timestamps-key", default="timestamps") + parser.add_argument("--velocities-key", default="joint_velocities") + parser.add_argument("--start-index", type=int, default=0) + parser.add_argument("--stop-index", type=int) + parser.add_argument("--eef-key") + parser.add_argument( + "--eef-description", + type=Path, + help="Explicit EefPoseSamples frame/TCP/unit description as JSON", + ) + parser.add_argument( + "--phases", + type=Path, + help="Optional SourcePhase declarations, indexed relative to the selected slice", + ) + args = parser.parse_args() + input_path = args.input.resolve() + output = args.output.resolve() + if output == input_path.parent or input_path.parent in output.parents: + parser.error("Output must be outside the original recording directory") + if bool(args.eef_key) != bool(args.eef_description): + parser.error("--eef-key and --eef-description must be supplied together") + with np.load(input_path, allow_pickle=False) as data: + total = len(data[args.positions_key]) + stop = total if args.stop_index is None else args.stop_index + if not 0 <= args.start_index < stop <= total: + parser.error("Invalid source slice") + section = slice(args.start_index, stop) + q = data[args.positions_key][section] + t = data[args.timestamps_key][section] + v = data[args.velocities_key][section] if args.velocities_key in data else None + eef = None + if args.eef_key: + description = json.loads(args.eef_description.read_text()) + for key in ("reference_to_base", "link_to_tcp"): + description[key] = np.asarray(description[key], dtype=float) + eef = EefPoseSamples(poses=data[args.eef_key][section], **description) + phases = ( + () + if args.phases is None + else tuple(SourcePhase(**item) for item in json.loads(args.phases.read_text())) + ) + digest = hashlib.sha256(input_path.read_bytes()).hexdigest() + raw = Real2SimTrajectory( + source_id="external_recording", + source_revision=digest, + unit_id=f"samples_{args.start_index}_{stop}", + robot_id=args.robot_id, + robot_revision=args.robot_revision, + joint_names=tuple(args.joint_names), + joint_units=tuple(args.joint_units), + positions=q, + timestamps=t, + time_unit=args.time_unit, + calibration_revision="unverified", + annotation_revision="unverified", + phases=phases, + velocities=v, + eef=eef, + provenance={ + "recording_sha256": digest, + "source_sample_range": [args.start_index, stop], + "scope": "Recorded channels only; no full robot or physical scene verified", + }, + ) + factors = { + "rad": ("rad", 1.0), + "deg": ("rad", np.pi / 180), + "m": ("m", 1.0), + "mm": ("m", 0.001), + } + layout = RobotJointLayout( + args.robot_id, + args.robot_revision, + tuple(args.joint_names), + tuple(factors[u][0] for u in args.joint_units), + q[0] * np.array([factors[u][1] for u in args.joint_units]), + ) + unit = normalize_real2sim( + raw, + layout, + Real2SimNormalizationCfg(control_dt=args.control_dt, allow_unverified=True), + ) + output.mkdir(parents=True, exist_ok=True) + for path in (output / "normalization.json", output / "trajectory.npz"): + if path.exists(): + raise FileExistsError( + f"{path} already exists; choose a fresh output directory" + ) + report = unit.report_dict() + report["example_scope"] = ( + "Offline diagnostic only; not a qualified Pick or committed dataset episode" + ) + (output / "normalization.json").write_text(json.dumps(report, indent=2) + "\n") + arrays = { + "raw_qpos": raw.positions, + "raw_timestamps": raw.timestamps, + "qpos": unit.positions, + "qvel": unit.velocities, + "dt": unit.dt, + "execution_times_s": unit.time_map.execution_times, + "source_query_times_s": unit.time_map.source_query_times, + } + arrays.update( + { + "map_" + name: getattr(unit.time_map, name) + for name in unit.time_map.__dataclass_fields__ + } + ) + if raw.velocities is not None: + arrays["raw_qvel"] = raw.velocities + if raw.eef is not None: + arrays["raw_eef"] = raw.eef.poses + arrays["eef_pose"] = unit.eef_poses + np.savez_compressed(output / "trajectory.npz", **arrays) + print( + json.dumps( + { + "source_samples": len(raw.positions), + "execution_samples": len(unit.positions), + "input_prerequisites_verified": unit.eligible, + "revision": unit.revision, + "output": str(output), + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/examples/real2sim/physics_preview.py b/examples/real2sim/physics_preview.py new file mode 100644 index 000000000..47298ebef --- /dev/null +++ b/examples/real2sim/physics_preview.py @@ -0,0 +1,962 @@ +# ---------------------------------------------------------------------------- +# 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 physical interaction preview for one external qpos recording. + +Uses the shared fixed-cadence playback host. Objects are initialized only at +reset and then move exclusively through gravity/contact; no attachment or +object-pose playback is used. Scene geometry and hardware calibration remain +estimates, so this diagnostic never grants real2sim source admission. +""" + +from __future__ import annotations + +import argparse +import json +import hashlib +import math +from pathlib import Path +import subprocess +import xml.etree.ElementTree as ET + +import cv2 +import numpy as np +import pytorch_kinematics as pk +import torch +from scipy.spatial.transform import Rotation + +from embodichain.data import get_data_path +from embodichain.data_pipeline.real2sim.gripper import ( + GripperPositionCalibration, + gripper_release_duration, +) +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import ( + RobotCfg, + URDFCfg, + RenderCfg, + LightCfg, + DefaultPhysicsCfg, + RigidObjectCfg, + JointDrivePropertiesCfg, + RigidBodyPhysicsCfg, + MassPropertiesCfg, + RigidBodyMaterialCfg, + CollisionPropertiesCfg, + DefaultRigidBodyPropertiesCfg, + ArticulationRootPropertiesCfg, +) +from embodichain.lab.sim.material import VisualMaterialCfg +from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.sensors import CameraCfg, ContactSensorCfg +from embodichain.lab.visualization import ( + add_viser_args_to_parser, + visualization_cfg_from_args, +) +from embodichain.lab.sim.motion.execution import ( + play_joint_trajectory, + JointTrajectoryPlaybackCfg, +) +from examples.real2sim.render_preview import ( + ARM_NAMES, + _build_paths, + _decode_video, + _text, +) + +__all__: list[str] = [] +CONTROL_DT = 0.01 +PHYSICS_DT = 0.001 + + +def _physics(friction: float, mass: float | None = None) -> RigidBodyPhysicsCfg: + return RigidBodyPhysicsCfg( + mass_props=None if mass is None else MassPropertiesCfg(mass=mass), + collision_props=CollisionPropertiesCfg(contact_offset=0.0005, rest_offset=0.0), + material_props=RigidBodyMaterialCfg( + static_friction=friction, dynamic_friction=friction, restitution=0.0 + ), + rigid_props=DefaultRigidBodyPropertiesCfg( + linear_damping=0.02, angular_damping=0.05, enable_ccd=True + ), + ) + + +def _model_geometry(arm_path: Path, hand_path: Path, q: np.ndarray, driver: float): + arm = pk.build_serial_chain_from_urdf(arm_path.read_bytes(), "fr3_link8").to( + dtype=torch.float64 + ) + flange = arm.forward_kinematics(torch.from_numpy(q[None])).get_matrix()[0].numpy() + hand = pk.build_chain_from_urdf(hand_path.read_bytes()).to(dtype=torch.float64) + joints = [ + j + for j in ET.parse(hand_path).getroot().findall("joint") + if j.get("type") != "fixed" + ] + values = { + j.get("name"): torch.tensor( + [ + driver + * ( + float(j.find("mimic").get("multiplier", "1")) + if j.find("mimic") is not None + else 1.0 + ) + ], + dtype=torch.float64, + ) + for j in joints + } + fk = hand.forward_kinematics(values) + pads = [] + for name in ("left_inner_finger", "right_inner_finger"): + pads.append( + (fk[name].get_matrix()[0].numpy() @ np.array([-0.0078, 0, 0.033, 1]))[:3] + ) + center = (flange @ np.r_[np.mean(pads, axis=0), 1])[:3] + yaw = math.atan2(flange[1, 0], flange[0, 0]) + return center, yaw, flange + + +def _box( + sim, uid, size, pos, color, *, yaw=0.0, dynamic=False, mass=None, friction=0.8 +): + return sim.add_rigid_object( + RigidObjectCfg( + uid=uid, + body_type="dynamic" if dynamic else "static", + init_pos=tuple(float(v) for v in pos), + init_rot=(0.0, 0.0, float(np.rad2deg(yaw))), + shape=CubeCfg( + size=list(size), + visual_material=VisualMaterialCfg( + uid=( + "physics_bowl_mat" if uid.startswith("bowl_") else uid + "_mat" + ), + base_color=[*color, 1.0], + roughness=0.75, + ), + ), + attrs=_physics(friction, mass), + ) + ) + + +class _ObservedPlayback: + """Observe the existing playback host without changing its control clock.""" + + def __init__( + self, + sim, + robot, + block, + sensor, + targets, + initial_z, + bowl_center, + table_z, + cameras, + writer, + fps, + label, + *, + open_driver_rad=0.0, + release_tolerance_rad=0.03, + ): + self.sim = sim + self.sim_config = sim.sim_config + self.robot = robot + self.block = block + self.sensor = sensor + self.targets = targets + self.initial_z = initial_z + self.bowl_center = bowl_center + self.table_z = table_z + self.cameras = cameras + self.writer = writer + self.fps = fps + self.label = label + self.open_driver_rad = open_driver_rad + self.release_tolerance_rad = release_tolerance_rad + self.samples = [] + self.time = 0.0 + self.frame_count = 0 + self.actor_names = {} + + def update(self, physics_dt, steps): + self.sim.update(physics_dt, steps, render_final_step=False) + self.time += physics_dt * steps + index = min(len(self.samples), len(self.targets) - 1) + pose = self.block.get_local_pose()[0].cpu().numpy().copy() + q = self.robot.get_qpos()[0].cpu().numpy().copy() + self.sensor.update() + contact = self.sensor.get_data() + valid = contact["is_valid"].cpu().numpy().astype(bool) + pairs = contact["user_ids"].cpu().numpy()[valid] + names = [] + for pair in pairs: + labels = [] + for actor in pair: + if actor >= 0: + actor = int(actor) + if actor not in self.actor_names: + info = self.sensor.get_actor_info(actor) + self.actor_names[actor] = info.path + ( + "/" + info.link_name if info.link_name else "" + ) + labels.append(self.actor_names[actor]) + names.extend(labels) + left = any("left_inner_finger" in x or "left_outer_finger" in x for x in names) + right = any( + "right_inner_finger" in x or "right_outer_finger" in x for x in names + ) + impulse = contact["impulse"].cpu().numpy()[valid] + target = self.targets[index] + active = self.active_ids + error = float(np.abs(q[active] - target).max()) + flange = self.robot.get_link_pose("fr3_link8")[0].cpu().numpy().copy() + self.samples.append( + { + "t": self.time, + "pose": pose, + "flange": flange, + "qpos": q, + "target": target.copy(), + "bilateral": left and right, + "left": left, + "right": right, + "contact_count": len(pairs), + "impulse": float(impulse.max()) if len(impulse) else 0.0, + "tracking_error": error, + } + ) + if self.writer is not None: + while self.frame_count / self.fps <= self.time + 1e-10: + with self.sim.render_frame(): + self.sim.sync_render_state() + images = [] + for camera in self.cameras: + camera.update() + images.append( + cv2.cvtColor( + camera.get_data()["color"][0, :, :, :3].cpu().numpy(), + cv2.COLOR_RGB2BGR, + ) + ) + canvas = np.zeros((450, 640, 3), np.uint8) + canvas[:] = (22, 28, 36) + canvas[40:400] = images[1] + canvas[46:158, 6:206] = cv2.resize(images[0], (200, 112)) + _text(canvas, self.label, (16, 27), scale=0.65) + height = (pose[2] - self.initial_z) * 1000 + _text( + canvas, + f"t={self.time:05.2f}s block lift {height:+.0f} mm", + (12, 421), + scale=0.55, + ) + _text( + canvas, + f"contact L/R: {int(left)}/{int(right)} q error: {error:.3f} rad", + (12, 444), + scale=0.48, + ) + self.writer.stdin.write(canvas.tobytes()) + self.frame_count += 1 + if len(self.samples) % 300 == 0: + print( + f"{self.label}: t={self.time:.2f}s lift={pose[2]-self.initial_z:.3f} contacts={left}/{right} err={error:.4f}", + flush=True, + ) + + +def _resting_checks(poses, times, block_size, bowl_center, table_z, bowl_radius): + positions = poses[:, :3] + final = positions[-1] + tail = poses[times >= times[-1] - 1.0] + signs = np.array([[x, y, z] for x in (-1, 1) for y in (-1, 1) for z in (-1, 1)]) + corners = (signs * np.array(block_size) / 2) @ Rotation.from_quat( + poses[-1, 3:] + ).as_matrix().T + final + inner_radius = bowl_radius - 0.004 + inside = bool( + ( + np.linalg.norm(corners[:, :2] - bowl_center[:2], axis=1) + <= inner_radius + 0.001 + ).all() + and corners[:, 2].min() >= table_z + 0.008 - 0.002 + and corners[:, 2].max() <= table_z + 0.06 + 0.002 + ) + speed = float( + np.linalg.norm(np.diff(tail[:, :3], axis=0) / CONTROL_DT, axis=1).max() + ) + rotations = Rotation.from_quat(tail[:, 3:]) + angle = float((rotations * rotations[-1].inv()).magnitude().max()) + angular_speed = float( + (rotations[1:] * rotations[:-1].inv()).magnitude().max() / CONTROL_DT + ) + stable = bool( + np.linalg.norm(tail[:, :3] - final, axis=1).max() < 0.005 + and speed < 0.03 + and angle < np.deg2rad(3) + and angular_speed < 0.2 + ) + return { + "final_inside_container": inside, + "stable_last_second": stable, + "max_final_second_object_speed_m_s": speed, + "max_final_second_rotation_rad": angle, + "max_final_second_angular_speed_rad_s": angular_speed, + } + + +def _summary(observer, block_size, bowl_radius): + data = observer.samples + positions = np.array([s["pose"][:3] for s in data]) + times = np.array([s["t"] for s in data]) + heights = positions[:, 2] - observer.initial_z + bilateral = np.array([s["bilateral"] for s in data]) + start = None + sustained = 0.0 + for timestamp, valid in zip(times, bilateral & (heights > 0.04)): + if valid: + if start is None: + start = timestamp + sustained = max(sustained, float(timestamp - start)) + else: + start = None + final = positions[-1] + resting = _resting_checks( + np.array([s["pose"] for s in data]), + times, + block_size, + observer.bowl_center, + observer.table_z, + bowl_radius, + ) + inside = resting["final_inside_container"] + stable = resting["stable_last_second"] + q = np.array([s["qpos"] for s in data])[:, observer.active_ids] + targets = np.array([s["target"] for s in data]) + measured_open = ( + np.abs(q[:, -1] - observer.open_driver_rad) <= observer.release_tolerance_rad + ) + released_for = gripper_release_duration(times, measured_open, bilateral) + release_pass = released_for >= 1.0 - 1e-9 + return { + "max_block_lift_m": float(heights.max()), + "continuous_bilateral_lift_s": sustained, + "grasp_lift_pass": sustained >= 1.0, + "final_inside_container": bool(inside), + "stable_last_second": bool(stable), + "continuous_release_s": released_for, + "release_pass": release_pass, + "physical_task_pass": bool( + sustained >= 1.0 and inside and stable and release_pass + ), + "max_joint_tracking_error_rad": float(max(s["tracking_error"] for s in data)), + "arm_tracking_rmse_rad": float( + np.sqrt(np.mean((q[:, :7] - targets[:, :7]) ** 2)) + ), + "max_arm_tracking_error_rad": float(np.abs(q[:, :7] - targets[:, :7]).max()), + **resting, + "final_block_xyz": final.tolist(), + "physical_validation_scope": "estimated_proxy_scene", + "real2sim_qualified": False, + } + + +def _compose(args, output, unit, results): + source_frames = _decode_video(args.source_video) + captures = [cv2.VideoCapture(str(output / f"candidate_{i}.mp4")) for i in range(3)] + counts = [int(c.get(cv2.CAP_PROP_FRAME_COUNT)) for c in captures] + if len(set(counts)) != 1: + raise RuntimeError("Candidate video clocks differ") + writer = subprocess.Popen( + [ + "ffmpeg", + "-v", + "error", + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "-s", + "1280x960", + "-r", + str(args.fps), + "-i", + "pipe:0", + "-an", + "-c:v", + "libx264", + "-preset", + "fast", + "-crf", + "20", + "-pix_fmt", + "yuv420p", + "-movflags", + "+faststart", + str(output / "physical_comparison.mp4"), + ], + stdin=subprocess.PIPE, + ) + try: + for frame in range(counts[0]): + time = frame / args.fps + tick = min(int(round(time / CONTROL_DT)), len(unit.positions) - 1) + index = int( + np.abs( + unit.source.seconds - unit.time_map.source_query_times[tick] + ).argmin() + ) + source = np.full((450, 640, 3), (22, 28, 36), np.uint8) + source[40:400] = source_frames[min(index, len(source_frames) - 1)] + _text(source, "REAL RECORDING / state-clock mapping", (16, 27), scale=0.62) + _text( + source, + f"t={min(time,float(unit.source.seconds[-1])):05.2f}s source frame {index}", + (12, 422), + scale=0.55, + ) + _text( + source, "Camera frame offset remains unverified", (12, 444), scale=0.48 + ) + panels = [] + for cap in captures: + ok, panel = cap.read() + if not ok: + raise RuntimeError("Candidate frame decode failed") + panels.append(panel) + canvas = np.full((960, 1280, 3), (18, 23, 30), np.uint8) + canvas[:450] = np.concatenate([source, panels[0]], axis=1) + canvas[450:900] = np.concatenate(panels[1:], axis=1) + phase = ( + "FREE MOTION VARIANTS" + if tick < unit.phases[0].stop + else "PROTECTED MOTION / PHYSICAL CONTACT" + ) + _text( + canvas, f"PHYSICS 1 ms | CONTROL 10 ms | {phase}", (16, 925), scale=0.57 + ) + _text( + canvas, + f"Estimated foreground + uncalibrated mount | measured arm qpos; explicit gripper preload {args.grip_preload:.3f} rad", + (16, 950), + (175, 195, 210), + 0.47, + ) + writer.stdin.write(canvas.tobytes()) + if frame in (0, 255, counts[0] - 1): + cv2.imwrite(str(output / f"comparison_{frame:04d}.jpg"), canvas) + card = np.full((960, 1280, 3), (18, 23, 30), np.uint8) + _text(card, "MEASURED PHYSICAL DIAGNOSTIC", (70, 120), scale=1.1) + _text( + card, + "Same scene, reset-only object initialization, no object attachment", + (70, 180), + scale=0.72, + ) + for i, name in enumerate( + ("Nominal", "Variant A / seed 7", "Variant B / seed 11") + ): + rows = [r for r in results if r["candidate"] == i] + passed = sum(r["physical_task_pass"] for r in rows) + _text( + card, + f"{name}: {passed}/{len(rows)} passed", + (80, 300 + 150 * i), + (225, 230, 235), + 0.95, + ) + _text( + card, + f"Lift {min(r['max_block_lift_m'] for r in rows)*1000:.0f} mm | hold {min(r['continuous_bilateral_lift_s'] for r in rows):.2f} s | released {min(r['continuous_release_s'] for r in rows):.2f} s | stable 1 s", + (80, 345 + 150 * i), + (160, 195, 215), + 0.67, + ) + _text( + card, + "Proxy geometry and model-based gripper mapping; real2sim calibration remains unverified.", + (70, 850), + (210, 190, 130), + 0.62, + ) + for _ in range(2 * args.fps): + writer.stdin.write(card.tobytes()) + finally: + for cap in captures: + cap.release() + writer.stdin.close() + if writer.wait() != 0: + raise RuntimeError("Comparison video encoding failed") + + +def _run(args): + output = args.output.resolve() + source = args.input.resolve() + if output == source.parent or source.parent in output.parents: + raise ValueError("Output must be outside the original recording directory") + output.mkdir(parents=True, exist_ok=False) + arm = Path(get_data_path("Franka/Panda/Panda.urdf")) + hand = Path(get_data_path("Robotiq2F85/Robotiq2F85.urdf")) + unit, arm_paths, _, flange_paths, protected, diagnostics, _ = _build_paths( + args, arm + ) + calibration = GripperPositionCalibration( + **json.loads(args.gripper_calibration.read_text()) + ) + open_driver = float( + np.interp( + calibration.max_width_m, + calibration.widths_m, + calibration.driver_positions_rad, + ) + ) + driver_range = abs( + calibration.driver_positions_rad[-1] - calibration.driver_positions_rad[0] + ) + if not 0 < args.release_open_tolerance < driver_range / 2: + raise ValueError( + "Release tolerance must be smaller than half the calibrated driver range" + ) + with np.load(source, allow_pickle=False) as raw: + scalar = raw["source_gripper_position"].copy() + if len(scalar) != len(unit.source.positions) or len(scalar) <= 160: + raise ValueError( + "This recording profile requires aligned source samples 105 and 160" + ) + driver = np.interp( + unit.time_map.source_query_times, + unit.source.seconds, + calibration.decode(scalar), + ) + # Measurement positions are not force commands: explicitly record a local + # squeeze preload while closure is nonzero; all three candidates share it. + driver_targets = driver + args.grip_preload * (driver > 0.03) + driver_targets = np.minimum(driver_targets, 0.725) + contact_index = int( + np.argmin(np.abs(unit.source.seconds - unit.source.seconds[105])) + ) + center, yaw, _ = _model_geometry( + arm, + hand, + unit.source.positions[contact_index], + float(calibration.decode(scalar)[105]), + ) + center[2] += args.base_height + width = calibration.max_width_m * (1 - float(scalar[105])) + block_size = (width, 0.05, 0.045) + table_z = float(center[2] - block_size[2] / 2) + bowl_center, _, _ = _model_geometry( + arm, hand, unit.source.positions[160], float(calibration.decode(scalar)[160]) + ) + bowl_center[2] = table_z + bowl_radius = 0.085 + joint_tree = [ + j for j in ET.parse(hand).getroot().findall("joint") if j.get("type") != "fixed" + ] + arm_stiff = {name: 1e4 for name in ARM_NAMES} + arm_damp = {name: 600.0 for name in ARM_NAMES} + arm_effort = {name: (87.0 if i < 4 else 12.0) for i, name in enumerate(ARM_NAMES)} + for j in joint_tree: + name = j.get("name") + active = name == "finger_joint" + arm_stiff[name] = args.grip_stiffness if active else 0.0 + arm_damp[name] = 8.0 if active else 0.0 + arm_effort[name] = args.grip_effort if active else 100.0 + assembly = URDFCfg( + components=[ + {"component_type": "arm", "urdf_path": str(arm), "transform": np.eye(4)}, + {"component_type": "hand", "urdf_path": str(hand), "transform": np.eye(4)}, + ], + fname="physics_franka_robotiq", + fpath_prefix=str(output / "robot"), + ) + cfg = RobotCfg( + uid="franka", + urdf_cfg=assembly, + init_pos=(0.0, 0.0, args.base_height), + control_parts={"arm": list(ARM_NAMES), "eef": ["finger_joint"]}, + root_props=ArticulationRootPropertiesCfg( + fixed_base=True, + self_collision_enabled=False, + min_position_iters=16, + min_velocity_iters=4, + ), + attrs=_physics(args.friction), + joint_drive_props=JointDrivePropertiesCfg( + drive_type="force", + target_mode="position", + stiffness=arm_stiff, + damping=arm_damp, + max_effort=arm_effort, + ), + ) + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + device="cpu", + num_envs=1, + physics_cfg=DefaultPhysicsCfg(physics_dt=PHYSICS_DT, enable_ccd=True), + render_cfg=RenderCfg( + renderer="no-render" if args.no_render else "hybrid", + tone_mapping_enabled=True, + ), + robot_ik_gizmo=None, + startup_summary="off", + visualization=visualization_cfg_from_args(args), + ) + ) + all_results = [] + try: + if not args.no_render: + sim.set_emission_light([0.92, 0.95, 1.0], 30.0) + sim.set_ground_plane_visibility(False) + sim.add_light( + LightCfg( + uid="key", + light_type="direction", + intensity=8.0, + direction=(-0.5, 0.4, -1.0), + ) + ) + _box( + sim, + "table", + (1.3, 1.0, 0.04), + (0.35, 0, table_z - 0.02), + (0.18, 0.30, 0.34), + friction=args.friction, + ) + _box( + sim, + "bowl_floor", + (0.18, 0.18, 0.008), + (bowl_center[0], bowl_center[1], table_z + 0.004), + (0.8, 0.86, 0.9), + ) + for i in range(24): + angle = 2 * np.pi * i / 24 + _box( + sim, + f"bowl_wall_{i}", + (0.008, 2 * bowl_radius * np.sin(np.pi / 24) + 0.001, 0.06), + ( + bowl_center[0] + bowl_radius * np.cos(angle), + bowl_center[1] + bowl_radius * np.sin(angle), + table_z + 0.03, + ), + (0.8, 0.86, 0.9), + yaw=angle, + ) + block = _box( + sim, + "block", + block_size, + center, + (0.95, 0.72, 0.05), + yaw=yaw, + dynamic=True, + mass=args.mass, + friction=args.friction, + ) + robot = sim.add_robot(cfg) + sensor = sim.add_sensor( + ContactSensorCfg( + uid="block_contacts", + rigid_uid_list=["block"], + filter_need_both_actor=False, + max_contacts_per_env=128, + ) + ) + cameras = [] + if not args.no_render: + for uid, eye, target, fx in [ + ("overview", (1.20, -1.45, 1.0), (0.32, 0, 0.30), 620.0), + ("contact_view", (0.86, -0.76, 0.52), (0.50, -0.035, 0.045), 620.0), + ]: + cameras.append( + sim.add_sensor( + CameraCfg( + uid=uid, + width=640, + height=360, + intrinsics=(fx, fx, 320.0, 180.0), + extrinsics=CameraCfg.ExtrinsicsCfg( + eye=(eye[0], eye[1], eye[2] + args.base_height), + target=( + target[0], + target[1], + target[2] + args.base_height, + ), + up=(0.0, 0.0, 1.0), + ), + ) + ) + ) + sim.prepare() + active_names = (*ARM_NAMES, "finger_joint") + active_ids = [robot.joint_names.index(n) for n in active_names] + target_paths = np.concatenate( + [ + arm_paths, + np.broadcast_to( + driver_targets[None, :, None], (3, len(driver_targets), 1) + ), + ], + axis=2, + ) + settle_count = 150 + target_paths = np.concatenate( + [target_paths, np.repeat(target_paths[:, -1:, :], settle_count, axis=1)], + axis=1, + ) + dt = np.r_[0.0, np.full(target_paths.shape[1] - 1, CONTROL_DT)] + settings = { + "source_revision": unit.revision, + "source_input_sha256": hashlib.sha256(source.read_bytes()).hexdigest(), + "normalization": unit.report_dict(), + "gripper_calibration": json.loads(args.gripper_calibration.read_text()), + "physics_dt": PHYSICS_DT, + "control_dt": CONTROL_DT, + "full_joint_names": robot.joint_names, + "active_joint_names": active_names, + "active_joint_ids": active_ids, + "geometry_estimated": True, + "robot_base_height": args.base_height, + "table_top_z": table_z, + "block_initial_xyz": center.tolist(), + "block_yaw": yaw, + "block_size": block_size, + "bowl_center": bowl_center.tolist(), + "bowl_radius": bowl_radius, + "mass": args.mass, + "friction": args.friction, + "grip_preload_rad": args.grip_preload, + "grip_effort_nm": args.grip_effort, + "grip_stiffness": args.grip_stiffness, + "object_state_writes": "reset only", + "attachments": False, + "robot_control": "PD joint targets through shared play_joint_trajectory", + "calibration_status": "unverified; estimated scene and hardware mapping", + "source_diagnostics": diagnostics, + "acceptance": { + "minimum_lift_m": 0.04, + "continuous_bilateral_hold_s": 1.0, + "final_stability_s": 1.0, + "minimum_trailing_release_s": 1.0, + "release_open_driver_rad": open_driver, + "release_open_tolerance_rad": args.release_open_tolerance, + "release_evidence": "measured open driver and no bilateral grasp contact", + "final_position_spread_m": 0.005, + "final_speed_limit_m_s": 0.03, + "final_angle_limit_rad": float(np.deg2rad(3)), + "final_angular_speed_rad_s": 0.2, + "containment": "all oriented box corners above container floor and inside rim", + }, + } + (output / "settings.json").write_text(json.dumps(settings, indent=2) + "\n") + np.savez_compressed( + output / "candidate_targets.npz", + targets=target_paths, + dt=dt, + raw_driver=driver, + command_driver=driver_targets, + protected_mask=protected, + ) + for candidate in args.cases: + for repeat in range(args.repeats): + sim.reset_objects_state() + initial = torch.zeros((1, len(robot.joint_names)), dtype=torch.float32) + initial[:, active_ids] = torch.from_numpy( + target_paths[candidate, 0] + ).float() + for joint in joint_tree: + mimic = joint.find("mimic") + if mimic is not None: + initial[ + :, robot.joint_names.index(joint.get("name")) + ] = driver_targets[0] * float( + mimic.get("multiplier", "1") + ) + float( + mimic.get("offset", "0") + ) + robot.set_state( + qpos=initial, + target_qpos=initial, + qvel=torch.zeros_like(initial), + target_qvel=torch.zeros_like(initial), + clear_dynamics=True, + ) + for _ in range(60): + sim.update(PHYSICS_DT, 10, render_final_step=False) + initial_z = float(block.get_local_pose()[0, 2]) + writer = None + label = ["NOMINAL", "VARIANT A / seed 7", "VARIANT B / seed 11"][ + candidate + ] + if not args.no_render and repeat == 0: + writer = subprocess.Popen( + [ + "ffmpeg", + "-v", + "error", + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "-s", + "640x450", + "-r", + str(args.fps), + "-i", + "pipe:0", + "-an", + "-c:v", + "libx264", + "-crf", + "20", + "-pix_fmt", + "yuv420p", + "-movflags", + "+faststart", + str(output / f"candidate_{candidate}.mp4"), + ], + stdin=subprocess.PIPE, + ) + observer = _ObservedPlayback( + sim, + robot, + block, + sensor, + target_paths[candidate], + initial_z, + bowl_center, + table_z, + cameras, + writer, + args.fps, + label, + open_driver_rad=open_driver, + release_tolerance_rad=args.release_open_tolerance, + ) + observer.active_ids = active_ids + try: + play_joint_trajectory( + observer, + robot, + positions=torch.from_numpy( + target_paths[candidate : candidate + 1] + ).float(), + dt=torch.from_numpy(dt[None]).float(), + joint_ids=active_ids, + cfg=JointTrajectoryPlaybackCfg(control_dt=CONTROL_DT), + ) + finally: + if writer is not None: + writer.stdin.close() + if writer.wait() != 0: + raise RuntimeError("Physical preview encoder failed") + result = _summary(observer, block_size, bowl_radius) + result.update(candidate=candidate, repeat=repeat, physics_stepped=True) + all_results.append(result) + np.savez_compressed( + output / f"trace_{candidate}_{repeat}.npz", + t=[s["t"] for s in observer.samples], + object_pose=[s["pose"] for s in observer.samples], + flange=[s["flange"] for s in observer.samples], + qpos=[s["qpos"] for s in observer.samples], + target=[s["target"] for s in observer.samples], + bilateral=[s["bilateral"] for s in observer.samples], + left=[s["left"] for s in observer.samples], + right=[s["right"] for s in observer.samples], + contacts=[s["contact_count"] for s in observer.samples], + impulse=[s["impulse"] for s in observer.samples], + ) + (output / "results.json").write_text( + json.dumps(all_results, indent=2) + "\n" + ) + (output / "contact_actors.json").write_text( + json.dumps(observer.actor_names, indent=2) + "\n" + ) + print(json.dumps(result), flush=True) + finally: + sim.destroy(exit_process=False) + del robot, block, sensor, cameras, sim + if ( + not args.no_render + and args.source_video is not None + and set(args.cases) == {0, 1, 2} + ): + _compose(args, output, unit, all_results) + + +def main() -> None: + """Run bounded B=1 physical comparisons using estimated foreground geometry.""" + p = argparse.ArgumentParser(description=__doc__) + add_viser_args_to_parser(p) + p.add_argument("--input", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--source-video", type=Path) + p.add_argument("--gripper-calibration", type=Path, required=True) + p.add_argument("--free-stop", type=int, default=60) + p.add_argument("--scale", type=float, default=0.025) + p.add_argument("--base-height", type=float, default=0.75) + p.add_argument("--mass", type=float, default=0.08) + p.add_argument("--friction", type=float, default=0.8) + p.add_argument("--grip-preload", type=float, default=0.015) + p.add_argument("--grip-effort", type=float, default=2.0) + p.add_argument("--grip-stiffness", type=float, default=600.0) + p.add_argument("--release-open-tolerance", type=float, default=0.03) + p.add_argument("--cases", type=int, nargs="+", choices=(0, 1, 2), default=[0, 1, 2]) + p.add_argument("--repeats", type=int, default=1) + p.add_argument("--fps", type=int, default=30) + p.add_argument("--no-render", action="store_true") + args = p.parse_args() + physical_values = ( + args.mass, + args.friction, + args.grip_preload, + args.grip_effort, + args.grip_stiffness, + args.base_height, + args.release_open_tolerance, + ) + if ( + not all(math.isfinite(value) for value in physical_values) + or args.repeats < 1 + or args.fps < 1 + or args.mass <= 0 + or args.friction < 0 + or args.grip_preload < 0 + or args.grip_effort <= 0 + or args.grip_stiffness <= 0 + or args.release_open_tolerance <= 0 + or args.base_height < 0.1 + ): + p.error("Invalid physical parameters") + try: + _run(args) + finally: + # _run has unwound its observer/camera/robot references before native cleanup. + SimulationManager.flush_cleanup_queue() + + +if __name__ == "__main__": + main() diff --git a/examples/real2sim/prepare_pick_sample.py b/examples/real2sim/prepare_pick_sample.py new file mode 100644 index 000000000..c759eb5fa --- /dev/null +++ b/examples/real2sim/prepare_pick_sample.py @@ -0,0 +1,292 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Audit one original DROID HDF5 and prepare a conservative Pick source bundle. + +This optional example never marks a model estimate as a measured calibration. +Its bundle remains diagnostic until mount/TCP and camera alignment are reviewed. +""" + +from __future__ import annotations + +import argparse +from dataclasses import asdict +import hashlib +import json +from pathlib import Path +import xml.etree.ElementTree as ET + +import h5py +import numpy as np +import pytorch_kinematics as pk +import torch +import trimesh + +from embodichain.data import get_data_path +from embodichain.data_pipeline.real2sim import ( + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, + MimicJoint, + EefPoseSamples, +) +from embodichain.data_pipeline.real2sim.bundle import write_real2sim_bundle +from embodichain.data_pipeline.real2sim.normalization import Real2SimNormalizationCfg +from embodichain.data_pipeline.real2sim.gripper import GripperPositionCalibration + +__all__: list[str] = [] +UPSTREAM = "https://github.com/droid-dataset/droid/blob/" # pin supplied on CLI + + +def _sha(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _model_calibration(hand_path: Path, evidence: str) -> GripperPositionCalibration: + root = ET.parse(hand_path).getroot() + joints = [j for j in root.findall("joint") if j.get("type") != "fixed"] + chain = pk.build_chain_from_urdf(hand_path.read_bytes()).to(dtype=torch.float64) + vertices = {} + for link in ("left_inner_finger", "right_inner_finger"): + collision = root.find(f"link[@name='{link}']/collision") + if collision.find("origin").get("xyz") != "0 0 0": + raise ValueError("Unexpected finger collision offset") + mesh = trimesh.load( + hand_path.parent / collision.find("geometry/mesh").get("filename"), + force="mesh", + ) + vertices[link] = mesh.vertices[mesh.vertices[:, 2] > 0.02] + angles = np.linspace(0.0, 0.725, 80) + widths = [] + for angle in angles: + values = { + j.get("name"): torch.tensor( + [ + angle + * ( + float(j.find("mimic").get("multiplier", "1")) + if j.find("mimic") is not None + else 1.0 + ) + ], + dtype=torch.float64, + ) + for j in joints + } + transforms = chain.forward_kinematics(values) + points = { + k: trimesh.transform_points(v, transforms[k].get_matrix()[0].numpy()) + for k, v in vertices.items() + } + widths.append( + float( + points["left_inner_finger"][:, 0].min() + - points["right_inner_finger"][:, 0].max() + ) + ) + if np.any(np.diff(widths) >= 0) or widths[-1] <= 0 or widths[0] >= 0.085: + raise ValueError( + "Unexpected Robotiq contact-pad geometry; review lookup extraction" + ) + # The nominal 85 mm hardware range extends beyond this mesh's open-pad span. + # Preserve the discrepancy explicitly; endpoint extensions are estimates. + return GripperPositionCalibration( + "normalized_closure", + 0.085, + tuple([0.0, *widths[::-1], 0.085]), + tuple([0.725, *angles[::-1], 0.0]), + "robotiq2f85-model-pad-lookup-v1", + evidence, + status="unverified", + ) + + +def main() -> None: + """Prepare one source slice and retain unresolved calibration prerequisites.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--hdf5", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--upstream-revision", required=True) + parser.add_argument("--close-start", type=int, required=True) + parser.add_argument("--lift-start", type=int, required=True) + parser.add_argument("--stop", type=int, required=True) + args = parser.parse_args() + source_path = args.hdf5.resolve() + output = args.output.resolve() + if output == source_path.parent or source_path.parent in output.parents: + parser.error("Output must be outside the original source directory") + if len(args.upstream_revision) != 40 or any( + c not in "0123456789abcdef" for c in args.upstream_revision + ): + parser.error("Pin the upstream source-code commit SHA") + output.mkdir(parents=True, exist_ok=False) + arm_path = Path(get_data_path("Franka/Panda/Panda.urdf")) + hand_path = Path(get_data_path("Robotiq2F85/Robotiq2F85.urdf")) + evidence = UPSTREAM + args.upstream_revision + "/droid/franka/robot.py" + calibration = _model_calibration(hand_path, evidence) + (output / "gripper_calibration.json").write_text( + json.dumps(asdict(calibration), indent=2) + "\n" + ) + with h5py.File(source_path, "r") as f: + q = f["joint_positions"][:] + times = f["timestamps"][:] + scalar = f["gripper_positions"][:] + poses = f["gripper_pose"][:] + arrays = { + "raw_gripper_scalar": scalar, + "raw_qpos": q, + "timestamps_ms": times, + "source_flange_pose": poses, + } + camera_metrics = {} + for camera in f["intrinsic"]: + intrinsic = f[f"intrinsic/{camera}"][:] + extrinsic = f[f"extrinsic/{camera}_optimized"][:] + arrays[f"intrinsic_{camera}"] = intrinsic + arrays[f"extrinsic_{camera}"] = extrinsic + camera_metrics[camera] = { + "rotation_determinant": float(np.linalg.det(extrinsic[:3, :3])), + "alignment_status": "unverified", + "direction": "base_to_camera candidate; independent reprojection pending", + } + pipeline_metadata = json.loads( + bytes(f["extrinsic/_pipeline_metadata"][:]).decode() + ) + if not 0 < args.close_start < args.lift_start < args.stop <= len(q): + parser.error("Invalid Pick phase boundaries") + if ( + np.any(scalar[: args.close_start] > 0.02) + or np.ptp(scalar[args.lift_start : args.stop]) > 0.02 + ): + raise ValueError("Phase assumptions disagree with the recorded gripper events") + driver = calibration.decode(scalar) + hand = ET.parse(hand_path).getroot() + moving = [j for j in hand.findall("joint") if j.get("type") != "fixed"] + arm_names = tuple(f"fr3_joint{i}" for i in range(1, 8)) + names = (*arm_names, *(j.get("name") for j in moving)) + mimics = tuple( + MimicJoint( + j.get("name"), + j.find("mimic").get("joint"), + float(j.find("mimic").get("multiplier", "1")), + float(j.find("mimic").get("offset", "0")), + ) + for j in moving + if j.find("mimic") is not None + ) + full = np.c_[ + q, + *[ + driver + * ( + float(j.find("mimic").get("multiplier", "1")) + if j.find("mimic") is not None + else 1.0 + ) + for j in moving + ], + ] + model_revision = hashlib.sha256( + (_sha(arm_path) + _sha(hand_path)).encode() + ).hexdigest() + source = Real2SimTrajectory( + "droid_original_hdf5", + _sha(source_path), + "pick_source_samples_0_" + str(args.stop), + "franka_robotiq2f85", + model_revision, + names, + ("rad",) * len(names), + full[: args.stop], + times[: args.stop], + "ms", + "partial-model-calibration-v1", + "manual-event-annotation-v1", + phases=( + SourcePhase("approach", 0, args.close_start, "free"), + SourcePhase("close", args.close_start, args.lift_start, "contact"), + SourcePhase("lift", args.lift_start, args.stop, "hold"), + ), + eef=EefPoseSamples( + poses[: args.stop], + "robot_base", + "fr3_link8", + "observed_flange", + np.eye(4), + np.eye(4), + "fk_derived", + ), + provenance={ + "gripper_mapping": asdict(calibration), + "original_scalar_sha256": hashlib.sha256(scalar.tobytes()).hexdigest(), + "annotation_note": "Manual conservative Pick phases supported by source gripper transition and lifting motion; contact time and camera offset remain unverified.", + "source_sample_range": [0, args.stop], + "source_pose_is_independent_measurement": False, + }, + ) + layout = RobotJointLayout( + source.robot_id, + model_revision, + names, + ("rad",) * len(names), + full[0], + arm_names, + mimics, + ) + path = write_real2sim_bundle( + source, + layout, + Real2SimNormalizationCfg(control_dt=0.01, allow_unverified=True), + output / "pick_bundle", + ) + chain = pk.build_serial_chain_from_urdf(arm_path.read_bytes(), "fr3_link8").to( + dtype=torch.float64 + ) + fk = chain.forward_kinematics(torch.from_numpy(q)).get_matrix().numpy() + rmse = float(np.sqrt(np.mean(np.sum((fk[:, :3, 3] - poses[:, :3]) ** 2, axis=1)))) + np.savez_compressed( + output / "calibration_evidence.npz", **arrays, model_driver_rad=driver + ) + report = { + "bundle": str(path.relative_to(output)), + "bundle_sha256": _sha(path), + "source_sha256": _sha(source_path), + "gripper_encoding": "normalized_closure", + "gripper_convention_reference": evidence, + "lookup_status": "model-derived estimate; hardware calibration pending", + "flange_position_consistency_rmse_m": rmse, + "independent_pose_accuracy_test": False, + "camera_candidates": camera_metrics, + "input_calibration_pipeline_metadata": pipeline_metadata, + "pick_source_segments": { + "approach": [0, args.close_start], + "close": [args.close_start, args.lift_start], + "lift": [args.lift_start, args.stop], + }, + "formal_admission": "rejected_until_reviewed", + "unresolved": [ + "real flange-to-gripper mount and TCP", + "hardware jaw-width/driver calibration", + "independent camera reprojection and N-1 video offset", + "target simulator full-joint order and initial scene registration", + ], + } + (output / "calibration_report.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/examples/real2sim/render_preview.py b/examples/real2sim/render_preview.py new file mode 100644 index 000000000..175e60352 --- /dev/null +++ b/examples/real2sim/render_preview.py @@ -0,0 +1,716 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Render one external qpos recording and two unqualified motion variants. + +This is a kinematic inspection example: no physics stepping, object attachment, +physical-success claim or EpisodeSink commit. The qualified Real2SimSourceAdapter +continues to reject unverified input. The example uses the public residual +operator directly on an explicitly diagnostic template. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +from pathlib import Path +import subprocess +import xml.etree.ElementTree as ET + +import cv2 +import numpy as np +import torch +import pytorch_kinematics as pk + +from embodichain.data import get_data_path +from embodichain.data_pipeline.real2sim import ( + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, +) +from embodichain.data_pipeline.real2sim.normalization import ( + Real2SimNormalizationCfg, + normalize_real2sim, +) +from embodichain.compute.trajectory import differentiate_positions +from embodichain.lab.visualization import ( + add_viser_args_to_parser, + visualization_cfg_from_args, +) +from embodichain.lab.sim.motion.expansion import TrajectoryPhase, TrajectoryTemplate +from embodichain.lab.sim.motion.expansion.operators import joint_residual + +__all__: list[str] = [] + +ARM_NAMES = tuple(f"fr3_joint{i}" for i in range(1, 8)) +COLORS = [(230, 210, 70), (65, 195, 255), (215, 125, 245)] # BGR + + +def _sha(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _text( + frame: np.ndarray, + message: str, + xy: tuple[int, int], + color=(235, 240, 245), + scale=0.55, +) -> None: + cv2.putText( + frame, message, xy, cv2.FONT_HERSHEY_SIMPLEX, scale, color, 1, cv2.LINE_AA + ) + + +def _decode_video(path: Path) -> np.ndarray: + probe = json.loads( + subprocess.check_output( + [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=width,height", + "-of", + "json", + str(path), + ] + ) + )["streams"][0] + if (probe["width"], probe["height"]) != (640, 360): + raise ValueError( + "This optional recording example expects a 640x360 source video" + ) + raw = subprocess.check_output( + [ + "ffmpeg", + "-v", + "error", + "-i", + str(path), + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "pipe:1", + ] + ) + return np.frombuffer(raw, np.uint8).reshape(-1, 360, 640, 3) + + +def _build_paths(args: argparse.Namespace, arm_path: Path): + with np.load(args.input, allow_pickle=False) as data: + q = data["joint_positions"].copy() + timestamps = data["timestamps"].copy() + observed_pose = ( + data["source_gripper_pose_xyzw"].copy() + if "source_gripper_pose_xyzw" in data + else None + ) + if not 5 <= args.free_stop < len(q): + raise ValueError( + "--free-stop must leave an interior free window and protected suffix" + ) + root = ET.parse(arm_path).getroot() + limits = np.array( + [ + [ + float(root.find(f"joint[@name='{name}']/limit").get(k)) + for k in ("lower", "upper") + ] + for name in ARM_NAMES + ] + ) + velocity_limits = np.array( + [ + float(root.find(f"joint[@name='{name}']/limit").get("velocity")) + for name in ARM_NAMES + ] + ) + raw = Real2SimTrajectory( + source_id="external_motion_preview", + source_revision=_sha(args.input), + unit_id=args.input.stem, + robot_id="Franka-arm", + robot_revision=_sha(arm_path), + joint_names=ARM_NAMES, + joint_units=("rad",) * 7, + positions=q, + timestamps=timestamps, + time_unit="ms", + calibration_revision="preview-unverified", + annotation_revision="manual-conservative-free-window-v1", + phases=( + SourcePhase("free", 0, args.free_stop, "free", ("joint_residual",)), + SourcePhase("protected_suffix", args.free_stop, len(q), "contact"), + ), + provenance={ + "scope": "kinematic diagnostic only", + "annotation_note": "Conservative free prefix reviewed from source images; suffix combines approach/contact/transport/release and is locked, not finely segmented.", + }, + ) + layout = RobotJointLayout( + "Franka-arm", _sha(arm_path), ARM_NAMES, ("rad",) * 7, q[0], ARM_NAMES + ) + unit = normalize_real2sim( + raw, layout, Real2SimNormalizationCfg(control_dt=0.01, allow_unverified=True) + ) + free = unit.phases[0] + # Protect two samples at each original endpoint for differentiation stencils. + template = TrajectoryTemplate( + source_id=raw.source_id, + source_revision=unit.revision, + template_id=raw.unit_id, + joint_names=ARM_NAMES, + positions=torch.from_numpy(unit.positions.copy()), + dt=torch.from_numpy(unit.dt.copy()), + controlled_joint_indices=tuple(range(7)), + phases=( + TrajectoryPhase( + "diagnostic_free", + free.start + 1, + free.stop - 1, + "free", + ("joint_residual",), + ), + ), + allowed_operators=("joint_residual",), + ) + paths = [template] + for seed in (7, 11): + paths.append( + joint_residual( + template, + joint_limits=torch.from_numpy(limits), + normalized_scale=args.scale, + generator=torch.Generator().manual_seed(seed), + ) + ) + qpos = np.stack([p.positions.numpy() for p in paths]) + qvel = np.stack( + [ + differentiate_positions(p.positions[None], p.dt[None])[0].numpy() + for p in paths + ] + ) + protected = np.ones(len(unit.positions), dtype=bool) + protected[free.start + 2 : free.stop - 2] = False + if not np.array_equal( + qpos[1:, protected], + np.broadcast_to(qpos[0, protected], qpos[1:, protected].shape), + ): + raise AssertionError("Expansion changed protected qpos") + if not np.array_equal( + qvel[1:, free.stop - 1 :], + np.broadcast_to(qvel[0, free.stop - 1 :], qvel[1:, free.stop - 1 :].shape), + ): + raise AssertionError("Expansion changed protected suffix velocity") + chain = pk.build_serial_chain_from_urdf(arm_path.read_text(), "fr3_link8").to( + dtype=torch.float64 + ) + if tuple(chain.get_joint_parameter_names()) != ARM_NAMES: + raise ValueError("Asset FK joint order differs from the source declaration") + fk = ( + chain.forward_kinematics(torch.from_numpy(qpos.reshape(-1, 7))) + .get_matrix() + .numpy() + .reshape(3, -1, 4, 4) + ) + source_fk = chain.forward_kinematics(torch.from_numpy(q)).get_matrix().numpy() + consistency = ( + None + if observed_pose is None + else float( + np.sqrt( + np.mean( + np.sum((source_fk[:, :3, 3] - observed_pose[:, :3]) ** 2, axis=1) + ) + ) + ) + ) + diagnostics = [] + for i in range(3): + diagnostics.append( + { + "name": ["nominal", "variant_a", "variant_b"][i], + "seed": [None, 7, 11][i], + "max_joint_deviation_rad": float(np.abs(qpos[i] - qpos[0]).max()), + "max_flange_deviation_m": float( + np.linalg.norm(fk[i, :, :3, 3] - fk[0, :, :3, 3], axis=1).max() + ), + "joint_limits_pass": bool( + ((qpos[i] >= limits[:, 0]) & (qpos[i] <= limits[:, 1])).all() + ), + "velocity_limits_pass": bool( + (np.abs(qvel[i]) <= velocity_limits).all() + ), + "physical_validation": "not_run", + "collision_validation": "not_run", + } + ) + if not all( + d["joint_limits_pass"] and d["velocity_limits_pass"] for d in diagnostics + ): + raise ValueError(f"Kinematic limit checks failed: {diagnostics}") + return unit, qpos, qvel, fk, protected, diagnostics, consistency + + +def _render(args: argparse.Namespace) -> None: + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.cfg import ( + RobotCfg, + URDFCfg, + RenderCfg, + LightCfg, + DefaultPhysicsCfg, + RigidObjectCfg, + ) + from embodichain.lab.sim.sensors import CameraCfg + from embodichain.lab.sim.shapes import CubeCfg + from embodichain.lab.sim.material import VisualMaterialCfg + + output = args.output.resolve() + source_dir = args.input.resolve().parent + if output == source_dir or source_dir in output.parents: + raise ValueError("Output must be outside the recording directory") + output.mkdir(parents=True, exist_ok=False) + arm_path = Path(get_data_path("Franka/Panda/Panda.urdf")) + hand_path = Path(get_data_path("Robotiq2F85/Robotiq2F85.urdf")) + unit, qpos, qvel, fk, protected, diagnostics, consistency = _build_paths( + args, arm_path + ) + gripper_calibration = None + gripper_driver = np.zeros(len(unit.positions)) + if args.gripper_calibration is not None: + from embodichain.data_pipeline.real2sim.gripper import ( + GripperPositionCalibration, + ) + + gripper_calibration = json.loads(args.gripper_calibration.read_text()) + calibration = GripperPositionCalibration(**gripper_calibration) + with np.load(args.input, allow_pickle=False) as recording: + scalar = recording["source_gripper_position"] + driver = calibration.decode(scalar) + gripper_driver = np.interp( + unit.time_map.source_query_times, unit.source.seconds, driver + ) + hand_joints = [ + j + for j in ET.parse(hand_path).getroot().findall("joint") + if j.get("type") != "fixed" + ] + frames = _decode_video(args.video) + if len(frames) not in (len(unit.source.positions), len(unit.source.positions) - 1): + raise ValueError( + "Expected N or N-1 source video frames; supply explicit synchronization for other layouts" + ) + assembly = URDFCfg( + components=[ + { + "component_type": "arm", + "urdf_path": str(arm_path), + "transform": np.eye(4), + }, + { + "component_type": "hand", + "urdf_path": str(hand_path), + "transform": np.eye(4), + }, + ], + fname="franka_robotiq", + fpath_prefix=str(output / "robot"), + ) + cfg = RobotCfg( + uid="preview_franka", + urdf_cfg=assembly, + control_parts={"arm": list(ARM_NAMES), "eef": ["finger_joint"]}, + ) + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + device="cpu", + num_envs=3, + arena_space=20.0, + physics_cfg=DefaultPhysicsCfg(), + render_cfg=RenderCfg(renderer="hybrid", tone_mapping_enabled=True), + startup_summary="off", + robot_ik_gizmo=None, + visualization=visualization_cfg_from_args(args), + ) + ) + writer = None + try: + sim.set_emission_light([0.92, 0.95, 1.0], 30.0) + sim.set_ground_plane_visibility(False) + sim.add_rigid_object( + RigidObjectCfg( + uid="preview_stage", + body_type="static", + init_pos=(0.0, 0.0, -0.025), + shape=CubeCfg( + size=[8.0, 8.0, 0.04], + visual_material=VisualMaterialCfg( + uid="preview_stage_material", + base_color=[0.12, 0.16, 0.20, 1.0], + roughness=0.85, + ), + ), + ) + ) + robot = sim.add_robot(cfg) + camera = sim.add_sensor( + CameraCfg( + uid="preview_camera", + width=640, + height=360, + far=4.0, + intrinsics=(620.0, 620.0, 320.0, 180.0), + extrinsics=CameraCfg.ExtrinsicsCfg( + eye=(1.20, -1.45, 1.00), + target=(0.32, 0.0, 0.34), + up=(0.0, 0.0, 1.0), + ), + ) + ) + sim.add_light( + LightCfg( + uid="key", + light_type="direction", + intensity=8.0, + direction=(-0.5, 0.4, -1.0), + ) + ) + sim.prepare() + ids = [robot.joint_names.index(n) for n in ARM_NAMES] + # Fixed-open gripper is explicit: the recorded scalar has no verified unit mapping. + full = torch.zeros( + (3, len(robot.joint_names)), dtype=torch.float32, device=robot.device + ) + duration = float(unit.time_map.execution_times[-1]) + render_times = np.arange(math.ceil(duration * args.fps) + 1) / args.fps + if args.max_frames: + render_times = render_times[: args.max_frames] + execution_indices = np.minimum( + np.rint(render_times / 0.01).astype(int), len(unit.positions) - 1 + ) + source_query = unit.time_map.source_query_times[execution_indices] + source_indices = np.abs( + unit.source.seconds[:, None] - source_query[None, :] + ).argmin(axis=0) + video_indices = np.clip( + source_indices + args.video_state_offset, 0, len(frames) - 1 + ) + np.savez_compressed( + output / "trajectories.npz", + qpos=qpos, + qvel=qvel, + flange_fk=fk, + protected_mask=protected, + dt=unit.dt, + gripper_driver_rad=gripper_driver, + render_times_s=render_times, + execution_indices=execution_indices, + source_indices=source_indices, + video_indices=video_indices, + **{ + "map_" + k: getattr(unit.time_map, k) + for k in unit.time_map.__dataclass_fields__ + }, + ) + report = { + "scope": "kinematic_motion_preview", + "physics_stepped": False, + "qualified_episode": False, + "normalization": unit.report_dict(), + "diagnostics": diagnostics, + "source_video_sha256": _sha(args.video), + "arm_asset_sha256": _sha(arm_path), + "gripper_asset_sha256": _sha(hand_path), + "assembled_asset": str(Path(assembly.fpath).relative_to(output)), + "assembled_asset_sha256": _sha(Path(assembly.fpath)), + "source_pose_vs_flange_position_rmse_m": consistency, + "source_pose_is_independent_measurement": False, + "free_source_samples": [0, args.free_stop], + "joint_residual_scale": args.scale, + "gripper_policy": ( + "fixed open" + if gripper_calibration is None + else "explicit width lookup; model estimate" + ), + "gripper_calibration": gripper_calibration, + "mount_policy": "identity flange-to-gripper mount for preview only", + "scene_policy": "neutral ground; no reconstructed objects or object animation", + "preview_camera": { + "intrinsics": camera.cfg.intrinsics, + "eye": camera.cfg.extrinsics.eye, + "target": camera.cfg.extrinsics.target, + "source_camera_calibrated": False, + }, + "video_state_offset": args.video_state_offset, + "video_alignment_verified": False, + "video_sync_policy": "state timestamp mapping; nearest frame; endpoint clamp for N-1 video", + "render_frames": len(render_times), + "fps": args.fps, + "rendered_duration_s": len(render_times) / args.fps, + "readback_max_error_rad": 0.0, + "flange_readback_max_error_m": 0.0, + } + writer = subprocess.Popen( + [ + "ffmpeg", + "-v", + "error", + "-f", + "rawvideo", + "-pix_fmt", + "bgr24", + "-s", + "1280x900", + "-r", + str(args.fps), + "-i", + "pipe:0", + "-an", + "-c:v", + "libx264", + "-preset", + "fast", + "-crf", + "20", + "-pix_fmt", + "yuv420p", + "-movflags", + "+faststart", + str(output / "comparison.mp4"), + ], + stdin=subprocess.PIPE, + ) + allxy = fk[:, :, :2, 3].reshape(-1, 2) + low, high = allxy.min(0) - 0.03, allxy.max(0) + 0.03 + + def plot_path(panel, index, tick): + origin = np.array([440, 290]) + span = np.array([160, -100]) + cv2.rectangle(panel, (430, 172), (630, 335), (26, 32, 40), -1) + _text(panel, "Flange path (XY)", (440, 193), scale=0.42) + for j in (0, index): + points = origin + (fk[j, :, :2, 3] - low) / (high - low) * span + cv2.polylines( + panel, + [points.astype(np.int32)], + False, + (95, 100, 110) if j == 0 else COLORS[index], + 1, + cv2.LINE_AA, + ) + now = origin + (fk[index, tick, :2, 3] - low) / (high - low) * span + cv2.circle(panel, tuple(now.astype(int)), 4, COLORS[index], -1) + _text( + panel, + f"delta {np.linalg.norm(fk[index,tick,:3,3]-fk[0,tick,:3,3])*1000:.1f} mm", + (440, 318), + COLORS[index], + 0.43, + ) + + for frame_number, tick in enumerate(execution_indices): + full[:, ids] = torch.from_numpy(qpos[:, tick]).to(full) + for joint in hand_joints: + mimic = joint.find("mimic") + multiplier = ( + 1.0 if mimic is None else float(mimic.get("multiplier", "1")) + ) + offset = 0.0 if mimic is None else float(mimic.get("offset", "0")) + full[:, robot.joint_names.index(joint.get("name"))] = ( + gripper_driver[tick] * multiplier + offset + ) + robot.set_qpos(full, target=False) + with sim.render_frame(): + # Camera.update() consumes native buffers; explicitly publish direct + # articulation writes without taking a positive-duration physics step. + sim.sync_render_state() + camera.update() + rendered = camera.get_data()["color"].cpu().numpy()[..., :3].copy() + error = float( + np.max(np.abs(robot.get_qpos().cpu().numpy()[:, ids] - qpos[:, tick])) + ) + report["readback_max_error_rad"] = max( + report["readback_max_error_rad"], error + ) + flange_error = float( + np.linalg.norm( + robot.get_link_pose("fr3_link8").cpu().numpy()[:, :3] + - fk[:, tick, :3, 3], + axis=1, + ).max() + ) + report["flange_readback_max_error_m"] = max( + report["flange_readback_max_error_m"], flange_error + ) + if error > 1e-5 or flange_error > 1e-4: + raise RuntimeError( + f"Preview state disagrees with source/FK: qpos={error}, flange={flange_error}" + ) + panels = [] + for i in range(4): + panel = np.full((410, 640, 3), (22, 28, 36), np.uint8) + picture = ( + frames[video_indices[frame_number]] + if i == 0 + else cv2.cvtColor(rendered[i - 1], cv2.COLOR_RGB2BGR) + ) + panel[38:398] = picture + title = [ + "REAL RECORDING / timestamp aligned", + "NOMINAL / recorded arm qpos", + "VARIANT A / seed 7", + "VARIANT B / seed 11", + ][i] + _text( + panel, + title, + (18, 26), + (235, 240, 245) if i == 0 else COLORS[i - 1], + 0.62, + ) + if i: + plot_path(panel, i - 1, tick) + _text( + panel, + "KINEMATIC | " + + ( + "fixed-open gripper" + if gripper_calibration is None + else "model-mapped gripper" + ), + (12, 385), + (30, 205, 250), + 0.5, + ) + panels.append(panel) + canvas = np.full((900, 1280, 3), (18, 23, 30), np.uint8) + canvas[:410] = np.concatenate(panels[:2], axis=1) + canvas[410:820] = np.concatenate(panels[2:], axis=1) + free_time = float(unit.time_map.execution_times[unit.phases[0].stop - 1]) + stage = ( + "FREE PREFIX / augmentation enabled" + if tick < unit.phases[0].stop + else "PROTECTED SUFFIX / identical arm motion" + ) + _text( + canvas, + f"t={render_times[frame_number]:05.2f}s | {stage}", + (20, 845), + (220, 230, 240), + 0.6, + ) + _text( + canvas, + "Motion inspection only: no grasp/contact validation; source video offset and mount remain uncalibrated.", + (20, 870), + (170, 185, 205), + 0.46, + ) + x0, x1 = 20, 1260 + cv2.line(canvas, (x0, 889), (x1, 889), (65, 75, 85), 5) + cv2.line( + canvas, + (x0, 889), + (int(x0 + (x1 - x0) * free_time / duration), 889), + COLORS[1], + 5, + ) + cv2.circle( + canvas, + ( + int( + x0 + + (x1 - x0) + * min(render_times[frame_number], duration) + / duration + ), + 889, + ), + 6, + (245, 245, 245), + -1, + ) + writer.stdin.write(canvas.tobytes()) + if frame_number in ( + 0, + min(60, len(execution_indices) - 1), + len(execution_indices) - 1, + ): + cv2.imwrite(str(output / f"frame_{frame_number:04d}.jpg"), canvas) + if frame_number % 90 == 0: + print(f"Rendered {frame_number+1}/{len(execution_indices)}", flush=True) + writer.stdin.close() + if writer.wait() != 0: + raise RuntimeError("Video encoder failed") + writer = None + (output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + print( + json.dumps({"output": str(output), "diagnostics": diagnostics}, indent=2), + flush=True, + ) + finally: + if writer is not None: + writer.stdin.close() + writer.wait() + sim.destroy(exit_process=False) + del robot, camera, sim + SimulationManager.flush_cleanup_queue() + + +def main() -> None: + """Parse one optional recording and render a non-qualifying preview.""" + parser = argparse.ArgumentParser(description=__doc__) + add_viser_args_to_parser(parser) + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--video", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--free-stop", + type=int, + required=True, + help="Reviewed exclusive source index; remaining motion is locked", + ) + parser.add_argument("--scale", type=float, default=0.025) + parser.add_argument("--fps", type=int, default=30) + parser.add_argument("--video-state-offset", type=int, choices=(-1, 0, 1), default=0) + parser.add_argument( + "--gripper-calibration", + type=Path, + help="Explicit scalar/width/driver lookup JSON; never qualifies physical calibration", + ) + parser.add_argument( + "--max-frames", + type=int, + default=0, + help="Bounded renderer smoke check; zero renders the whole recording", + ) + args = parser.parse_args() + if args.fps <= 0 or args.max_frames < 0: + parser.error("fps must be positive and max-frames nonnegative") + _render(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark/trajectory/anchored_batch.py b/scripts/benchmark/trajectory/anchored_batch.py new file mode 100644 index 000000000..50c27f67d --- /dev/null +++ b/scripts/benchmark/trajectory/anchored_batch.py @@ -0,0 +1,298 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Compare anchored NumPy, Torch and Warp batch resampling with accuracy checks. + +Run: python -m scripts.benchmark.trajectory.anchored_batch --batches 1 13 64 128 +""" + +from __future__ import annotations + +import argparse +from datetime import datetime +from pathlib import Path +import platform +import time +from collections.abc import Callable + +import numpy as np +import psutil +import torch +import warp as wp + +from embodichain.compute.trajectory import ( + prepare_anchored_batch, + resample_with_anchors, + resample_with_anchors_batched, +) + + +def _fixture(batch: int, samples: int, dtype: torch.dtype) -> tuple: + """Generate seeded ragged paths with independent event counts and padding.""" + rng = np.random.default_rng(192) + lengths = rng.integers(max(8, samples // 2), samples + 1, size=batch) + q = torch.full((batch, samples, 13), float("nan"), dtype=dtype) + times = torch.full((batch, samples), float("nan"), dtype=torch.float64) + anchors = torch.full((batch, 6), -1, dtype=torch.long) + counts = torch.tensor(rng.integers(0, 7, size=batch)) + for b, n in enumerate(lengths): + clock = np.r_[0.0, np.cumsum(rng.uniform(0.05, 0.09, n - 1))] + times[b, :n] = torch.tensor(clock) + q[b, :n] = torch.tensor( + np.sin(clock[:, None] * np.linspace(0.1, 1.0, 13)[None]) + b * 0.01, + dtype=dtype, + ) + anchors[b, : counts[b]] = torch.tensor( + rng.choice(n, int(counts[b]), replace=False) + ) + return ( + q, + times, + anchors, + dict(source_lengths=torch.tensor(lengths), anchor_counts=counts), + ) + + +def _numpy_batch( + q: torch.Tensor, times: torch.Tensor, anchors: torch.Tensor, metadata: dict +) -> np.ndarray: + """Evaluate individual references and apply identical terminal padding.""" + rows = [] + for b, n in enumerate(metadata["source_lengths"].tolist()): + k = int(metadata["anchor_counts"][b]) + row, _ = resample_with_anchors( + q[b, :n].numpy(), times[b, :n].numpy(), 0.01, anchors[b, :k].tolist() + ) + rows.append(row) + out = np.empty( + (len(rows), max(len(r) for r in rows), q.shape[-1]), dtype=rows[0].dtype + ) + for b, row in enumerate(rows): + out[b] = row[-1] + out[b, : len(row)] = row + return out + + +def _measure(fn: Callable, device: str, repeats: int) -> tuple[object, dict]: + """Measure synchronized warmed wall time, process RSS and Torch allocations.""" + fn() + if device.startswith("cuda"): + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + process = psutil.Process() + rss = process.memory_info().rss + gpu = torch.cuda.memory_allocated() if device.startswith("cuda") else 0 + durations = [] + result = None + for _ in range(repeats): + if device.startswith("cuda"): + torch.cuda.synchronize() + start = time.perf_counter() + result = fn() + if device.startswith("cuda"): + torch.cuda.synchronize() + durations.append((time.perf_counter() - start) * 1000) + current_gpu = torch.cuda.memory_allocated() if device.startswith("cuda") else 0 + peak = torch.cuda.max_memory_allocated() if device.startswith("cuda") else 0 + return result, dict( + cost_time_ms=float(np.median(durations)), + cpu_delta_mb=(process.memory_info().rss - rss) / 1024**2, + gpu_delta_mb=(current_gpu - gpu) / 1024**2, + peak_gpu_mb=peak / 1024**2, + ) + + +def _table(rows: list[dict]) -> list[str]: + """Format one rectangular Markdown table.""" + keys = list(rows[0]) + return [ + "| " + " | ".join(keys) + " |", + "| " + " | ".join(["---"] * len(keys)) + " |", + *( + "| " + + " | ".join( + f"{r[k]:.6g}" if isinstance(r[k], float) else str(r[k]) for k in keys + ) + + " |" + for r in rows + ), + ] + + +def run_all_benchmarks() -> None: + """Run backend/device comparisons and write exactly one Markdown report.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batches", type=int, nargs="+", default=[1, 13, 64, 128]) + parser.add_argument("--samples", type=int, default=192) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--threads", type=int, default=4) + parser.add_argument("--dtype", choices=["float32", "float64"], default="float64") + parser.add_argument("--output", type=Path) + args = parser.parse_args() + if ( + args.samples < 8 + or args.repeats < 1 + or args.threads < 1 + or any(b < 1 for b in args.batches) + ): + parser.error("Need samples >=8 and positive batches, repeats and threads") + torch.set_num_threads(args.threads) + dtype = getattr(torch, args.dtype) + perf = [] + metrics = [] + methods = [ + ("numpy_single", "cpu", None), + ("torch_cpu", "cpu", "torch"), + ("warp_cpu", "cpu", "warp"), + ] + if torch.cuda.is_available(): + methods.extend( + [("torch_cuda", "cuda:0", "torch"), ("warp_cuda", "cuda:0", "warp")] + ) + else: + print("CUDA unavailable: skipping GPU comparisons") + for batch in args.batches: + q, times, anchors, metadata = _fixture(batch, args.samples, dtype) + expected = _numpy_batch(q, times, anchors, metadata) + for name, device, backend in methods: + if backend is None: + apply = lambda: _numpy_batch(q, times, anchors, metadata) + prepare_ms = 0.0 + plan = None + end_to_end = apply + else: + device_q = q.to(device) + plan, preparation = _measure( + lambda: prepare_anchored_batch( + times, 0.01, anchors, device=device, **metadata + ), + device, + args.repeats, + ) + prepare_ms = preparation["cost_time_ms"] + apply = lambda: plan.apply(device_q, backend=backend) + + # Deliberately include input/output transfer and all mapping + # preparation here to model the offline normalization caller. + def end_to_end() -> torch.Tensor: + """Include CPU input/output transfer and clock preparation.""" + out, _ = resample_with_anchors_batched( + q.to(device), times, 0.01, anchors, backend=backend, **metadata + ) + return out.cpu() + + out, cost = _measure(apply, device, args.repeats) + _, total = _measure(end_to_end, device, args.repeats) + actual = ( + out.detach().cpu().numpy() if isinstance(out, torch.Tensor) else out + ) + error = float(np.max(np.abs(actual - expected))) + tolerance = 2e-7 if dtype == torch.float32 else 2e-14 + passed = np.all( + np.isclose(actual, expected, atol=tolerance, rtol=tolerance), + axis=(1, 2), + ) + exact = True + if plan is not None: + source = q.numpy() + bit_type = np.int32 if dtype == torch.float32 else np.int64 + for b, row in enumerate(plan.rows): + exact &= np.array_equal( + actual[b, row.anchor_execution_indices].view(bit_type), + source[b, row.anchor_source_indices].view(bit_type), + ) + if not passed.all() or not exact: + raise AssertionError( + f"{name}: batch {batch} failed accuracy/anchor check" + ) + perf.append( + dict( + batch=batch, + impl=name, + **cost, + prepare_ms=prepare_ms, + end_to_end_ms=total["cost_time_ms"], + ) + ) + metrics.append( + dict( + batch=batch, + impl=name, + success_rate=float(passed.mean()), + max_abs_error=error, + anchors_exact=exact, + ) + ) + print( + f"B={batch:4d} {name:13s} apply={cost['cost_time_ms']:9.3f} ms total={total['cost_time_ms']:9.3f} ms max_error={error:.3g}" + ) + leaderboard = [] + for name, _, _ in methods: + rows = [r for r in metrics if r["impl"] == name] + costs = [r["cost_time_ms"] for r in perf if r["impl"] == name] + leaderboard.append( + dict( + algorithm=name, + overall_success_rate=float(np.mean([r["success_rate"] for r in rows])), + mean_apply_ms=float(np.mean(costs)), + ) + ) + leaderboard.sort(key=lambda r: (-r["overall_success_rate"], r["mean_apply_ms"])) + leaderboard = [dict(rank=i + 1, **row) for i, row in enumerate(leaderboard)] + report = ( + args.output + or Path("outputs/benchmarks") + / f"anchored_batch_{datetime.now():%Y%m%d_%H%M%S}.md" + ) + report.parent.mkdir(parents=True, exist_ok=True) + gpu_name = ( + torch.cuda.get_device_name(0) if torch.cuda.is_available() else "unavailable" + ) + lines = [ + "# Anchored batch resampling benchmark", + "", + f"Generated: {datetime.now().isoformat(timespec='seconds')}", + f"Platform: {platform.platform()}; CPU: {platform.processor()}; GPU: {gpu_name}", + f"Torch {torch.__version__}; Warp {wp.__version__}; dtype {args.dtype}; CPU threads {args.threads}; repeats {args.repeats}", + f"Batches {args.batches}; N_max {args.samples}; D=13; control_dt=0.01 s; seed=192.", + "", + "## Time & Memory", + "", + *_table(perf), + "", + "## Success & Other Metrics", + "", + *_table(metrics), + "", + "## Leaderboard", + "", + *_table(leaderboard), + "", + "- Success denotes numerical agreement with individual NumPy references, not physical task success.", + "- apply includes input validation; prepared batch plans and device inputs are reused. NumPy single includes preparation each call.", + "- end_to_end includes mapping preparation, CPU input transfer and CPU output transfer; warm compilation is excluded everywhere.", + "- prepare_ms is recorded separately; zero for NumPy means no separately prepared implementation.", + "- CPU memory is process RSS delta; GPU fields are Torch allocator delta/absolute peak around apply and exclude driver/context/JIT memory.", + "- Shorter paths use terminal padding; exact anchors, valid lengths and row independence are also covered by the production tests.", + "", + ] + with report.open("x", encoding="utf-8") as handle: + handle.write("\n".join(lines)) + print(f"Markdown report saved: {report}") + + +if __name__ == "__main__": + run_all_benchmarks() diff --git a/tests/compute/test_trajectory_anchors.py b/tests/compute/test_trajectory_anchors.py new file mode 100644 index 000000000..2783fa11d --- /dev/null +++ b/tests/compute/test_trajectory_anchors.py @@ -0,0 +1,67 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Hard-event and timing-provenance tests for source-neutral resampling.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from embodichain.compute.trajectory import resample_with_anchors + + +def test_preserves_anchor_positions_and_bidirectional_mapping(): + t = np.array([0.0, 0.011, 0.021, 0.041, 0.070]) + q = np.c_[np.sin(t * 20), t**2] + result, mapping = resample_with_anchors(q, t, 0.01, [2, 3]) + np.testing.assert_array_equal( + result[mapping.anchor_execution_indices], q[mapping.anchor_source_indices] + ) + np.testing.assert_allclose(np.diff(mapping.execution_times), 0.01) + reconstructed = q[mapping.left_indices] + mapping.weights[:, None] * ( + q[mapping.right_indices] - q[mapping.left_indices] + ) + np.testing.assert_allclose(result, reconstructed) + np.testing.assert_allclose( + mapping.source_query_times[mapping.anchor_execution_indices], + t[mapping.anchor_source_indices], + ) + assert mapping.execution_times[-1] >= t[-1] + with pytest.raises(ValueError): + mapping.weights.setflags(write=True) + + +def test_sub_tick_anchors_fail_instead_of_erasing_contact(): + t = np.array([0.0, 0.011, 0.012, 0.04]) + with pytest.raises(ValueError, match="anchors cannot fit"): + resample_with_anchors(np.zeros((4, 2)), t, 0.01, [1, 2]) + + +@pytest.mark.parametrize( + "times", + [[0.0, 0.1, 0.1], [0.0, 0.2, 0.1], [0.1, 0.2, 0.3], [0.0, float("nan"), 0.3]], +) +def test_invalid_source_clock_is_rejected(times): + with pytest.raises(ValueError): + resample_with_anchors(np.zeros((3, 2)), np.array(times), 0.01, []) + + +def test_allocation_bound_is_checked_before_grid_creation(): + with pytest.raises(ValueError, match="max_samples"): + resample_with_anchors( + np.zeros((2, 1)), np.array([0.0, 1000.0]), 1e-6, [], max_samples=100 + ) diff --git a/tests/compute/test_trajectory_anchors_batch.py b/tests/compute/test_trajectory_anchors_batch.py new file mode 100644 index 000000000..f49de359d --- /dev/null +++ b/tests/compute/test_trajectory_anchors_batch.py @@ -0,0 +1,385 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Ragged batch, backend parity and exact-event contracts for anchored resampling.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from embodichain.compute.trajectory import ( + AnchoredBatchPlan, + AnchoredTimeMap, + prepare_anchored_batch, + resample_with_anchors, + resample_with_anchors_batched, +) + +_EXECUTIONS = [ + ("cpu", "torch"), + ("cpu", "warp"), + pytest.param("cuda:0", "torch", marks=pytest.mark.gpu), + pytest.param("cuda:0", "warp", marks=pytest.mark.gpu), +] + + +def _inputs(dtype=torch.float64, device="cpu"): + if str(device).startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + times = torch.full((3, 6), float("nan"), dtype=torch.float64) + q = torch.full((3, 6, 4), float("nan"), dtype=dtype) + clocks = [ + [0.0, 0.011, 0.021, 0.041, 0.070, 0.097], + [0.0, 0.025, 0.047, 0.083], + [0.0, 0.09], + ] + for b, clock in enumerate(clocks): + t = torch.tensor(clock, dtype=torch.float64) + times[b, : len(t)] = t + q[b, : len(t)] = torch.stack( + (torch.sin(t * 20) + b, t * t, -t, t * 0 + 3), dim=1 + ).to(dtype) + q[0, 2, 3] = -0.0 + return ( + q.to(device), + times, + torch.tensor([[4, 2, 2, -999], [2, 1, -999, -999], [-999] * 4]), + dict( + source_lengths=torch.tensor([6, 4, 2]), + anchor_counts=torch.tensor([3, 2, 0]), + ), + ) + + +def _assert_reference(result, plan, q, times, anchors, args): + host = result.detach().cpu().numpy() + source = q.detach().cpu().numpy() + for b, row in enumerate(plan.rows): + n = args["source_lengths"][b].item() + k = args["anchor_counts"][b].item() + expected, mapping = resample_with_anchors( + source[b, :n], times[b, :n].numpy(), 0.01, anchors[b, :k].tolist() + ) + m = len(expected) + tolerance = 2e-7 if q.dtype == torch.float32 else 2e-14 + np.testing.assert_allclose( + host[b, :m], expected, atol=tolerance, rtol=tolerance + ) + for field in mapping.__dataclass_fields__: + np.testing.assert_array_equal(getattr(row, field), getattr(mapping, field)) + # Check the representation too: copying anchors must preserve signed zero. + bits = np.int32 if q.dtype == torch.float32 else np.int64 + np.testing.assert_array_equal( + host[b, row.anchor_execution_indices].view(bits), + source[b, row.anchor_source_indices].view(bits), + ) + np.testing.assert_array_equal( + host[b, m:], np.broadcast_to(source[b, n - 1], host[b, m:].shape) + ) + assert plan.valid_counts[b] == m + assert plan.valid_mask[b].sum() == m + assert torch.isfinite(result).all() + + +@pytest.mark.parametrize("device,backend", _EXECUTIONS) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_ragged_batch_matches_single_paths_and_preserves_exact_anchors( + device, backend, dtype +): + q, times, anchors, args = _inputs(dtype, device) + original = q.clone() + result, plan = resample_with_anchors_batched( + q, times, 0.01, anchors, backend=backend, **args + ) + _assert_reference(result, plan, q, times, anchors, args) + assert result.device == q.device and result.dtype == dtype + assert plan.valid_counts.device == q.device + torch.testing.assert_close(q, original, rtol=0, atol=0, equal_nan=True) + # Returned metadata is owned, and mutations cannot change a reusable plan. + counts = plan.valid_counts + counts.zero_() + assert (plan.valid_counts >= 2).all() + + +@pytest.mark.parametrize("device,backend", _EXECUTIONS) +def test_reuses_mapping_for_changed_values_and_noncontiguous_paths(device, backend): + q, times, anchors, args = _inputs(device=device) + plan = prepare_anchored_batch(times, 0.01, anchors, device=device, **args) + storage = torch.stack((q, q), dim=-1) + changed = storage[..., 0] + assert not changed.is_contiguous() + changed.add_(0.125) + result = plan.apply(changed, backend=backend) + _assert_reference(result, plan, changed, times, anchors, args) + fresh, _ = resample_with_anchors_batched( + changed, times, 0.01, anchors, backend=backend, **args + ) + torch.testing.assert_close(result, fresh, rtol=0, atol=0) + + +@pytest.mark.parametrize("device,backend", _EXECUTIONS) +def test_random_batches_match_independent_scalar_interpolation(device, backend): + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + rng = np.random.default_rng(192) + batch, width, dimensions = 13, 41, 7 + lengths = rng.integers(10, width + 1, size=batch) + times = np.full((batch, width), np.nan) + values = np.full((batch, width, dimensions), np.nan) + anchors = np.full((batch, 4), -1) + for b, n in enumerate(lengths): + times[b, :n] = np.r_[0.0, np.cumsum(rng.uniform(0.015, 0.09, n - 1))] + values[b, :n] = rng.normal(size=(n, dimensions)) + anchors[b] = rng.choice(n, size=4, replace=False) + q = torch.tensor(values, device=device) + result, plan = resample_with_anchors_batched( + q, + torch.tensor(times), + 0.01, + torch.tensor(anchors), + source_lengths=torch.tensor(lengths), + anchor_counts=torch.full((batch,), 4), + backend=backend, + ) + for b, row in enumerate(plan.rows): + n, m = lengths[b], len(row.execution_times) + expected = np.stack( + [ + np.interp(row.source_query_times, times[b, :n], values[b, :n, d]) + for d in range(dimensions) + ], + axis=1, + ) + np.testing.assert_allclose( + result[b, :m].cpu().numpy(), expected, atol=2e-14, rtol=2e-14 + ) + + +def test_known_event_tick_mapping_and_stationary_trajectory(): + q = torch.tensor( + [[[2.0], [5.0], [8.0]], [[-1.0], [-1.0], [-1.0]]], dtype=torch.float64 + ) + times = torch.tensor([[0.0, 0.021, 0.070], [0.0, 0.03, 0.04]], dtype=torch.float64) + out, plan = resample_with_anchors_batched( + q, + times, + 0.01, + torch.tensor([[1], [1]]), + source_lengths=torch.tensor([3, 3]), + anchor_counts=torch.tensor([1, 1]), + ) + np.testing.assert_array_equal(plan.rows[0].anchor_execution_indices, [0, 3, 7]) + np.testing.assert_allclose( + out[0, :, 0], [2, 3, 4, 5, 5.75, 6.5, 7.25, 8], atol=1e-14 + ) + assert torch.equal(out[1], torch.full_like(out[1], -1)) + assert plan.valid_counts.tolist() == [8, 5] + + +@pytest.mark.parametrize( + "bad_field", ["lengths", "counts", "anchors", "clock", "finite"] +) +def test_invalid_valid_prefixes_fail_with_no_implicit_row_dropping(bad_field): + q, times, anchors, args = _inputs() + if bad_field == "lengths": + args["source_lengths"][1] = 7 + elif bad_field == "counts": + args["anchor_counts"][1] = 5 + elif bad_field == "anchors": + anchors[1, 0] = 4 + elif bad_field == "clock": + times[1, 1] = times[1, 0] + else: + q[1, 1] = float("nan") + with pytest.raises(ValueError): + resample_with_anchors_batched(q, times, 0.01, anchors, **args) + + +def test_dense_events_report_the_failing_row(): + q = torch.zeros((2, 4, 2), dtype=torch.float64) + times = torch.tensor( + [[0.0, 0.011, 0.021, 0.04], [0.0, 0.011, 0.012, 0.04]], dtype=torch.float64 + ) + with pytest.raises(ValueError, match="Batch row 1:.*anchors cannot fit"): + resample_with_anchors_batched( + q, + times, + 0.01, + torch.tensor([[1, 2], [1, 2]]), + source_lengths=torch.tensor([4, 4]), + anchor_counts=torch.tensor([2, 2]), + ) + + +def test_total_padding_and_output_element_bounds_are_checked(): + q, times, anchors, args = _inputs() + with pytest.raises(ValueError, match="max_samples"): + resample_with_anchors_batched(q, times, 0.01, anchors, max_samples=20, **args) + plan = prepare_anchored_batch(times, 0.01, anchors, **args) + with pytest.raises(ValueError, match="max_output_elements"): + plan.apply(q, max_output_elements=10) + + +@pytest.mark.parametrize( + "field", + [ + "source_lengths", + "anchor_counts", + "anchor_indices", + "positions", + "source_times", + "shape", + ], +) +def test_invalid_shapes_and_dtypes_are_rejected(field): + q, times, anchors, args = _inputs() + if field in args: + args[field] = args[field].float() + elif field == "anchor_indices": + anchors = anchors.float() + elif field == "positions": + q = q.to(torch.float16) + elif field == "source_times": + times = times.long() + else: + q = q[:2] + with pytest.raises(ValueError): + resample_with_anchors_batched(q, times, 0.01, anchors, **args) + + +def test_zero_explicit_anchors_preserve_source_endpoints(): + q, times, _, args = _inputs() + out, plan = resample_with_anchors_batched( + q, + times, + 0.01, + torch.empty((3, 0), dtype=torch.long), + source_lengths=args["source_lengths"], + anchor_counts=torch.zeros(3, dtype=torch.long), + ) + assert all(len(row.anchor_source_indices) == 2 for row in plan.rows) + assert torch.equal(out[:, 0], q[:, 0]) + + +def test_explicit_backend_errors_propagate_and_torch_autograd_remains_available( + monkeypatch, +): + from embodichain.compute.trajectory._warp import anchored + + q, times, anchors, args = _inputs() + plan = prepare_anchored_batch(times, 0.01, anchors, **args) + + def fail(*args): + raise RuntimeError("test launch failed") + + monkeypatch.setattr(anchored, "evaluate", fail) + with pytest.raises(RuntimeError, match="test launch failed"): + plan.apply(q, backend="warp") + with pytest.raises(ValueError, match="backend"): + plan.apply(q, backend="unknown") + q.requires_grad_() + with pytest.raises(ValueError, match="autograd"): + plan.apply(q, backend="warp") + plan.apply(q).sum().backward() + assert q.grad is not None and torch.isfinite(q.grad).all() + for b, n in enumerate(args["source_lengths"]): + assert torch.count_nonzero(q.grad[b, n:]) == 0 + + +@pytest.mark.gpu +def test_warp_follows_nondefault_torch_stream(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + import warp as wp + + wp.init() + q, times, anchors, args = _inputs(device="cuda:0") + plan = prepare_anchored_batch(times, 0.01, anchors, device=q.device, **args) + # Warm compilation separately, then order writes, launch and consumption on + # one non-default Torch stream without relying on default-stream ordering. + plan.apply(q, backend="warp") + torch.cuda.synchronize() + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + changed = q + 0.25 + result = plan.apply(changed, backend="warp") + consumed = result + 0.5 + stream.synchronize() + torch.testing.assert_close( + consumed, plan.apply(q + 0.25) + 0.5, rtol=1e-14, atol=1e-14 + ) + + +def test_batch_row_permutation_does_not_change_individual_results(): + q, times, anchors, args = _inputs() + out, _ = resample_with_anchors_batched(q, times, 0.01, anchors, **args) + order = torch.tensor([2, 0, 1]) + reordered, plan = resample_with_anchors_batched( + q[order], + times[order], + 0.01, + anchors[order], + source_lengths=args["source_lengths"][order], + anchor_counts=args["anchor_counts"][order], + ) + torch.testing.assert_close(reordered, out[order], rtol=0, atol=0) + assert plan.source_lengths.tolist() == [2, 6, 4] + + +@pytest.mark.gpu +def test_cuda_metadata_preparation_and_plan_device_mismatch(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + q, times, anchors, args = _inputs(device="cuda:0") + plan = prepare_anchored_batch( + times.cuda(), + 0.01, + anchors.cuda(), + **{key: value.cuda() for key, value in args.items()}, + ) + assert plan.device == q.device + _assert_reference(plan.apply(q, backend="warp"), plan, q, times, anchors, args) + with pytest.raises(ValueError, match="plan device"): + plan.apply(q.cpu()) + + +@pytest.mark.parametrize("device,backend", _EXECUTIONS) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_direct_plan_construction_copies_noncanonical_anchor_brackets( + device, backend, dtype +): + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + clock = np.array([0.0, 1.0, 2.0]) + mapping = AnchoredTimeMap( + clock, + clock, + clock, + clock, + np.array([0, 0, 1]), + np.array([0, 1, 2]), + np.array([0.0, 1.0, 1.0]), + np.array([0, 1, 2]), + np.array([0, 1, 2]), + ) + positions = torch.tensor([[[2.0], [0.1], [-0.0]]], device=device, dtype=dtype) + result = AnchoredBatchPlan((mapping,), device=device).apply( + positions, backend=backend + ) + bit_dtype = torch.int32 if dtype == torch.float32 else torch.int64 + assert torch.equal(result.view(bit_dtype), positions.view(bit_dtype)) diff --git a/tests/data_pipeline/real2sim/test_bundle_gripper.py b/tests/data_pipeline/real2sim/test_bundle_gripper.py new file mode 100644 index 000000000..8814127ac --- /dev/null +++ b/tests/data_pipeline/real2sim/test_bundle_gripper.py @@ -0,0 +1,184 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Source bundle admission and explicit gripper calibration contracts.""" + +from __future__ import annotations +from dataclasses import replace +import hashlib +import json +import numpy as np +import pytest +from embodichain.data_pipeline.real2sim import ( + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, +) +from embodichain.data_pipeline.real2sim.normalization import Real2SimNormalizationCfg +from embodichain.data_pipeline.real2sim.bundle import ( + load_real2sim_bundle, + write_real2sim_bundle, +) +from embodichain.data_pipeline.real2sim.gripper import ( + GripperPositionCalibration, + gripper_release_duration, +) + + +def _bundle(tmp_path, verified=True): + source = Real2SimTrajectory( + "fixture", + "raw1", + "pick0", + "robot", + "robot1", + ("arm", "finger"), + ("rad", "rad"), + np.zeros((12, 2)), + np.arange(12) * 0.01, + "s", + "cal1", + "ann1", + "verified" if verified else "unverified", + "verified", + ( + SourcePhase("approach", 0, 6, "free", ("joint_residual",)), + SourcePhase("close", 6, 9, "contact"), + SourcePhase("lift", 9, 12, "hold"), + ), + ) + layout = RobotJointLayout( + "robot", "robot1", source.joint_names, source.joint_units, np.zeros(2), ("arm",) + ) + return write_real2sim_bundle( + source, + layout, + Real2SimNormalizationCfg(control_dt=0.01, allow_unverified=not verified), + tmp_path / "bundle", + ) + + +def test_bundle_round_trip_and_data_hash_rejection(tmp_path): + path = _bundle(tmp_path) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + unit = load_real2sim_bundle(path, expected_sha256=digest) + assert unit.eligible and unit.phases[1].kind == "contact" + data = path.parent / "source.npz" + data.write_bytes(data.read_bytes() + b"corruption") + with pytest.raises(ValueError, match="array digest"): + load_real2sim_bundle(path, expected_sha256=digest) + + +def test_bundle_diagnostic_input_is_never_formally_admitted(tmp_path): + path = _bundle(tmp_path, verified=False) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + assert not load_real2sim_bundle( + path, expected_sha256=digest, require_verified=False + ).eligible + with pytest.raises(ValueError, match="unverified"): + load_real2sim_bundle(path, expected_sha256=digest) + + +def test_bundle_rejects_unknown_declarations_and_sidecar_traversal(tmp_path): + path = _bundle(tmp_path) + data = json.loads(path.read_text()) + data["arrays"] = "../source.npz" + path.write_text(json.dumps(data)) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + with pytest.raises(ValueError, match="sidecar path"): + load_real2sim_bundle(path, expected_sha256=digest) + data["arrays"] = "source.npz" + data["executable"] = "arbitrary.import" + path.write_text(json.dumps(data)) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + with pytest.raises(ValueError, match="exactly"): + load_real2sim_bundle(path, expected_sha256=digest) + + +def test_declared_gripper_encodings_have_consistent_physical_widths(): + calibration = GripperPositionCalibration( + "normalized_closure", + 0.08, + (0.0, 0.04, 0.08), + (0.7, 0.4, 0.0), + "fixture1", + "measured lookup", + status="verified", + ) + np.testing.assert_array_equal( + calibration.decode(np.array([0.0, 0.5, 1.0])), [0.0, 0.4, 0.7] + ) + np.testing.assert_array_equal( + replace(calibration, encoding="width_m").decode(np.array([0.08, 0.04, 0.0])), + [0.0, 0.4, 0.7], + ) + with pytest.raises(ValueError, match="domain"): + calibration.decode(np.array([1.1])) + with pytest.raises(ValueError, match="explicitly"): + replace(calibration, encoding="guess") + with pytest.raises(ValueError, match="Rejected"): + replace(calibration, status="rejected").decode(np.array([0.5])) + + +@pytest.mark.parametrize( + "opened,contact,expected", + [ + ([True] * 5, [False] * 5, 2.0), + ([False] * 5, [False] * 5, 0.0), + ([True] * 5, [True] * 5, 0.0), + ([True] * 5, [False, False, True, False, False], 0.5), + ([False, False, True, True, True], [False] * 5, 1.0), + ([True] * 5, [False, False, False, False, True], 0.0), + ], +) +def test_release_duration_requires_measured_opening_and_absence_of_grasp( + opened, contact, expected +): + times = np.arange(5) * 0.5 + assert ( + gripper_release_duration(times, np.array(opened), np.array(contact)) == expected + ) + + +def test_release_duration_uses_elapsed_time_instead_of_frame_count(): + times = np.array([0.0, 0.01, 0.05, 0.08]) + assert ( + gripper_release_duration(times, np.ones(4, dtype=bool), np.zeros(4, dtype=bool)) + == 0.08 + ) + for n in (0, 1): + assert ( + gripper_release_duration( + times[:n], np.ones(n, dtype=bool), np.zeros(n, dtype=bool) + ) + == 0.0 + ) + + +@pytest.mark.parametrize( + "times,opened,contact", + [ + ([0.0, 0.0], [True, True], [False, False]), + ([0.0, float("nan")], [True, True], [False, False]), + ([0.0, 1.0], [True], [False, False]), + ([0.0, 1.0], [1.0, 1.0], [False, False]), + ], +) +def test_release_duration_rejects_invalid_clocks_or_nonboolean_evidence( + times, opened, contact +): + with pytest.raises(ValueError): + gripper_release_duration(np.array(times), np.array(opened), np.array(contact)) diff --git a/tests/data_pipeline/real2sim/test_normalization.py b/tests/data_pipeline/real2sim/test_normalization.py new file mode 100644 index 000000000..d9b9741d8 --- /dev/null +++ b/tests/data_pipeline/real2sim/test_normalization.py @@ -0,0 +1,250 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Unit/frame/time normalization without DROID files or experimental code.""" + +from __future__ import annotations + +from dataclasses import replace +import hashlib +import json + +import numpy as np +import pytest + +from embodichain.data_pipeline.real2sim import ( + EefPoseSamples, + MimicJoint, + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, +) +from embodichain.data_pipeline.real2sim.normalization import ( + Real2SimNormalizationCfg, + normalize_real2sim, +) + + +def source(): + return Real2SimTrajectory( + source_id="fixture", + source_revision="raw-v1", + unit_id="pick-0", + robot_id="test_robot", + robot_revision="urdf-v1", + joint_names=("slide", "angle"), + joint_units=("mm", "deg"), + positions=np.array( + [[0.0, 0.0], [1.0, 5.0], [2.0, 10.0], [3.0, 15.0], [4.0, 20.0]] + ), + timestamps=np.array([1000, 1071, 1142, 1210, 1281]), + time_unit="ms", + calibration_revision="cal-v1", + annotation_revision="phase-v1", + calibration_status="verified", + annotation_status="verified", + phases=( + SourcePhase("approach", 0, 3, "free", ("joint_residual",)), + SourcePhase("contact", 3, 5, "contact"), + ), + ) + + +def layout(): + return RobotJointLayout( + "test_robot", + "urdf-v1", + ("angle", "slide", "follower", "unused"), + ("rad", "m", "rad", "rad"), + np.array([0.0, 0.0, 0.2, 0.4]), + ("angle",), + (MimicJoint("follower", "angle", -1.0, 0.2),), + ) + + +def test_mapping_units_mimics_hold_and_clock_are_explicit(): + raw = source() + target = layout() + unit = normalize_real2sim( + raw, + target, + Real2SimNormalizationCfg(control_dt=0.01, hold_joint_names=("unused",)), + ) + np.testing.assert_allclose( + unit.positions[-1], [np.deg2rad(20), 0.004, 0.2 - np.deg2rad(20), 0.4] + ) + np.testing.assert_array_equal(unit.positions[:, 2], 0.2 - unit.positions[:, 0]) + np.testing.assert_allclose( + unit.velocities[:, 2], -unit.velocities[:, 0], atol=1e-12 + ) + np.testing.assert_array_equal(unit.positions[:, 3], 0.4) + assert unit.eligible + assert unit.report["held_positions"]["unused"] == 0.4 + assert unit.phases[1].kind == "contact" + assert unit.phases[1].stop - unit.phases[1].start > 1 + assert unit.time_map.execution_times[-1] == pytest.approx(0.29) + assert ( + unit.revision + == normalize_real2sim( + raw, + target, + Real2SimNormalizationCfg(control_dt=0.01, hold_joint_names=("unused",)), + ).revision + ) + + +def test_missing_joint_requires_hold_and_mimic_cannot_be_expanded(): + with pytest.raises(ValueError, match="Missing independent joints"): + normalize_real2sim( + source(), layout(), Real2SimNormalizationCfg(control_dt=0.01) + ) + with pytest.raises(ValueError, match="independent"): + replace(layout(), expandable_joint_names=("follower",)) + + +def test_gripper_encoding_cannot_be_guessed(): + with pytest.raises(ValueError, match="explicit supported unit"): + replace(source(), joint_units=("mm", "normalized")) + + +def test_immutable_data_and_revision_changes(): + raw = source() + before = raw.fingerprint() + with pytest.raises(ValueError): + raw.positions.setflags(write=True) + assert before != replace(raw, annotation_revision="phase-v2").fingerprint() + cfg = Real2SimNormalizationCfg(control_dt=0.01, hold_joint_names=("unused",)) + a = normalize_real2sim(raw, layout(), cfg) + b = normalize_real2sim( + raw, replace(layout(), expandable_joint_names=("slide",)), cfg + ) + assert a.revision != b.revision + with pytest.raises(TypeError): + a.report["held_positions"]["unused"] = 1.0 + + +def test_unverified_data_is_diagnostic_only_and_rejected_data_always_fails(): + raw = replace(source(), calibration_status="unverified") + cfg = Real2SimNormalizationCfg(control_dt=0.01, hold_joint_names=("unused",)) + with pytest.raises(ValueError, match="diagnostic"): + normalize_real2sim(raw, layout(), cfg) + cfg.allow_unverified = True + assert not normalize_real2sim(raw, layout(), cfg).eligible + with pytest.raises(ValueError, match="Rejected"): + normalize_real2sim(replace(raw, calibration_status="rejected"), layout(), cfg) + + +def test_source_pose_conversion_does_not_claim_independent_measurement(): + raw = source() + poses = np.tile([1000.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0], (5, 1)) + reference = np.eye(4) + reference[1, 3] = 2.0 + tcp = np.eye(4) + tcp[2, 3] = 0.15 + eef = EefPoseSamples( + poses, "camera", "flange", "tcp", reference, tcp, "fk_derived", "wxyz", "mm" + ) + unit = normalize_real2sim( + replace(raw, eef=eef), + layout(), + Real2SimNormalizationCfg(control_dt=0.01, hold_joint_names=("unused",)), + ) + np.testing.assert_allclose( + unit.eef_poses[:, :3], np.tile([1.0, 2.0, 0.15], (len(unit.positions), 1)) + ) + assert unit.report["independent_eef_measurement"] is False + description = unit.report_dict()["source_description"]["eef"] + assert description["reference_frame"] == "camera" + assert description["quaternion_order"] == "wxyz" + assert description["position_unit"] == "mm" + np.testing.assert_array_equal(description["link_to_tcp"], tcp) + np.testing.assert_array_equal(description["reference_to_base"], reference) + + +def test_report_preserves_source_declarations_and_array_descriptors(): + raw = replace( + source(), provenance={"corrections": [{"reason": "unit declaration"}]} + ) + unit = normalize_real2sim( + raw, + layout(), + Real2SimNormalizationCfg(control_dt=0.01, hold_joint_names=("unused",)), + ) + report = json.loads(json.dumps(unit.report_dict())) + assert report["source_revision"] == raw.source_revision + assert report["source_description"]["joint_names"] == ["slide", "angle"] + assert report["source_description"]["joint_units"] == ["mm", "deg"] + assert report["source_description"]["time_unit"] == "ms" + assert report["source_description"]["provenance"]["corrections"][0] == { + "reason": "unit declaration" + } + for name, descriptor in report["time_mapping_arrays"].items(): + array = getattr(unit.time_map, name) + assert descriptor == { + "key": "map_" + name, + "shape": list(array.shape), + "dtype": array.dtype.str, + "sha256": hashlib.sha256(array.tobytes()).hexdigest(), + } + + +def test_nan_duplicate_names_duplicate_clock_and_mimic_cycles_fail(): + with pytest.raises(ValueError): + replace(source(), joint_names=("angle", "angle")) + with pytest.raises(ValueError): + replace(source(), timestamps=np.array([0, 1, 1, 2, 3])) + with pytest.raises(ValueError): + replace(source(), positions=np.full((5, 2), np.nan)) + with pytest.raises(ValueError, match="cycle"): + replace( + layout(), + mimics=(MimicJoint("angle", "follower"), MimicJoint("follower", "angle")), + expandable_joint_names=(), + ) + + +def test_initial_state_and_robot_revision_mismatch_fail(): + cfg = Real2SimNormalizationCfg(control_dt=0.01, hold_joint_names=("unused",)) + with pytest.raises(ValueError, match="same robot"): + normalize_real2sim(source(), replace(layout(), robot_revision="other"), cfg) + with pytest.raises(ValueError, match="initial positions"): + normalize_real2sim( + source(), + replace(layout(), initial_positions=np.array([0.1, 0, 0.1, 0.4])), + cfg, + ) + + +def test_integer_nanosecond_epoch_is_subtracted_before_conversion(): + raw = replace( + source(), + time_unit="ns", + timestamps=1_700_000_000_000_000_000 + + np.array([0, 71, 142, 210, 281], dtype=np.int64) * 1_000_000, + ) + np.testing.assert_allclose( + raw.seconds, [0, 0.071, 0.142, 0.210, 0.281], rtol=0, atol=1e-15 + ) + + +@pytest.mark.parametrize("authority", ["eef_pose", "hybrid"]) +def test_other_authorities_require_a_real_planning_provider(authority): + with pytest.raises(NotImplementedError, match="IK/replanning"): + normalize_real2sim( + source(), + layout(), + Real2SimNormalizationCfg(control_dt=0.01, authority=authority), + ) diff --git a/tests/gym/envs/task_program/test_catalog.py b/tests/gym/envs/task_program/test_catalog.py index 5431268ad..d38afe4eb 100644 --- a/tests/gym/envs/task_program/test_catalog.py +++ b/tests/gym/envs/task_program/test_catalog.py @@ -1244,3 +1244,51 @@ def __init__(self, **kwargs: object) -> None: registration ), ) + + +@dataclass(frozen=True) +class _ExpertSourceFactory: + provider_id: ClassVar[str] = "catalog_source" + revision: ClassVar[str] = "1" + capabilities: ClassVar[frozenset[str]] = frozenset( + {"replay", "phase_annotation", "replanning"} + ) + data_revision: str = "data1" + + def create(self, **kwargs): + raise AssertionError("Pure catalog preflight must not create a provider") + + +def test_expert_sources_are_fingerprinted_and_exact_call_occurrences_preflight(): + from embodichain.lab.task_program.integrations.expert_sources import ( + ExpertSourceRegistry, + ExpertSourceSelection, + ) + + base = _registration() + registry = ExpertSourceRegistry( + (_ExpertSourceFactory(),), + (ExpertSourceSelection("pick", "catalog_source", "1", 0),), + ) + registered = replace(base, expert_sources=registry) + changed = replace( + base, + expert_sources=ExpertSourceRegistry( + (_ExpertSourceFactory("data2"),), registry.selections + ), + ) + assert registered.fingerprint != base.fingerprint != changed.fingerprint + assert registered.fingerprint != changed.fingerprint + registered.assert_unchanged() + # Existing configured program helper exercises the ordinary preflight path. + program = decode_task_program(_program_payload()) + registered.catalog.preflight(program) + missing = replace( + base, + expert_sources=ExpertSourceRegistry( + registry.factories, + (ExpertSourceSelection("pick", "catalog_source", "1", 999),), + ), + ) + with pytest.raises(ValueError, match="occurrence|expanded program call"): + missing.catalog.preflight(program) diff --git a/tests/gym/envs/task_program/test_configured_integration.py b/tests/gym/envs/task_program/test_configured_integration.py index 98d435cdb..d1cb21e3e 100644 --- a/tests/gym/envs/task_program/test_configured_integration.py +++ b/tests/gym/envs/task_program/test_configured_integration.py @@ -1594,3 +1594,50 @@ def test_segmented_examples_preserve_overall_and_segment_instructions( for segment in segments: assert isinstance(segment.instruction, str) and segment.instruction.strip() assert segment.instruction != overall + + +def test_expert_source_paths_resolve_from_integration_file(tmp_path, monkeypatch): + """Source selection is data owned by integration, independent of process cwd.""" + import yaml + + task_name = "rubiks_cube_pick_place" + original = _config_path(task_name) + physical = _physical_components(original, _gym_config(task_name)) + selected = physical.config["task_program"] + program, task, _ = _resolve_task_program_components( + selected, base_dir=original.parent + ) + task = deepcopy(dict(task)) + task.setdefault("runtime_services", {})["expert_source_providers"] = [ + { + "call_id": "pick", + "call_index": 0, + "provider_id": "real2sim_qpos", + "revision": "1", + "target_profile_id": "fixture", + "target_asset_sha256": "1" * 64, + "bundle_path": "sources/bundle.json", + "bundle_sha256": "0" * 64, + } + ] + (tmp_path / "sources").mkdir() + (tmp_path / "sources/bundle.json").write_text("{}") + integration_path = tmp_path / "integration.yaml" + integration_path.write_text(yaml.safe_dump(task)) + policy = Path(selected["execution_policy"]) + if not policy.is_absolute(): + policy = (original.parent / policy).resolve() + elsewhere = tmp_path / "elsewhere" + elsewhere.mkdir() + monkeypatch.chdir(elsewhere) + _, resolved, _ = _resolve_task_program_components( + { + "program": str(program), + "integration": str(integration_path), + "execution_policy": str(policy), + }, + base_dir=tmp_path, + ) + assert resolved["runtime_services"]["expert_source_providers"][0][ + "bundle_path" + ] == str((tmp_path / "sources/bundle.json").resolve()) diff --git a/tests/gym/envs/task_program/test_environment.py b/tests/gym/envs/task_program/test_environment.py index 17c2d4690..f73650d6d 100644 --- a/tests/gym/envs/task_program/test_environment.py +++ b/tests/gym/envs/task_program/test_environment.py @@ -915,3 +915,131 @@ def test_factory_protocol_is_required() -> None: """Loose objects cannot enter the production assembly boundary.""" with pytest.raises(TypeError, match="TaskProgramEnvironmentFactory"): TaskProgramEnvironmentAdapter(object(), step_dt=_STEP_DT) + + +def test_one_pick_real2sim_source_runs_through_task_program_bridge( + tmp_path, monkeypatch +): + """Admitted synthetic data exercises B=1 source -> Skill -> bridge lifecycle.""" + from dataclasses import replace + import hashlib + import numpy as np + from embodichain.data_pipeline.real2sim import ( + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, + ) + from embodichain.data_pipeline.real2sim.normalization import ( + Real2SimNormalizationCfg, + ) + from embodichain.data_pipeline.real2sim.bundle import write_real2sim_bundle + from embodichain.lab.task_program.integrations.expert_sources import ( + ExpertSourceRegistry, + ExpertSourceSelection, + ) + from embodichain.lab.task_program.integrations.real2sim import ( + Real2SimExpertSourceProviderFactory, + ) + from embodichain.lab.sim.atomic_actions import RecoveryPolicy + + monkeypatch.setitem(globals(), "_BATCH_SIZE", 1) + preset = SkillPolicyPreset( + "safe", + action_option_templates={"pick": PickUpOptions()}, + effect_assurance=EffectAssurance.PROJECTED, + recovery_policy=RecoveryPolicy(max_replans=0, max_action_retries=0), + ) + profile = replace(_robot_profile(effect_monitors={}), presets={"safe": preset}) + asset = tmp_path / "robot.urdf" + asset.write_text("") + q = np.zeros((12, 2)) + q[:, 0] = np.linspace(0, 0.1, 12) + q[6:, 1] = 0.2 + raw = Real2SimTrajectory( + "synthetic", + "data1", + "pick0", + "fixture", + "r1", + ("arm_joint", "finger"), + ("rad", "rad"), + q, + np.arange(12) * _STEP_DT, + "s", + "cal1", + "phase1", + "verified", + "verified", + ( + SourcePhase("approach", 0, 6, "free", ("joint_residual",)), + SourcePhase("close", 6, 9, "contact"), + SourcePhase("lift", 9, 12, "hold"), + ), + ) + path = write_real2sim_bundle( + raw, + RobotJointLayout( + "fixture", "r1", raw.joint_names, raw.joint_units, q[0], ("arm_joint",) + ), + Real2SimNormalizationCfg(control_dt=_STEP_DT), + tmp_path / "source", + ) + registry = ExpertSourceRegistry( + ( + Real2SimExpertSourceProviderFactory( + str(path), + hashlib.sha256(path.read_bytes()).hexdigest(), + "fake_robot", + hashlib.sha256(asset.read_bytes()).hexdigest(), + ), + ), + (ExpertSourceSelection("pick", "real2sim_qpos", "1"),), + ) + + class Factory(_FakeEnvironmentFactory): + def create_robot_skill_profile(self): + return profile + + def create_atomic_action_engine(self, selected_profile): + engine = _engine(selected_profile) + engine.robot.joint_names = list(raw.joint_names) + engine.robot.cfg.fpath = str(asset) + engine.robot.compute_fk.side_effect = lambda qpos, **kwargs: torch.eye( + 4 + ).repeat(qpos.shape[0], 1, 1) + engine.robot.get_qpos_limits.return_value = torch.tensor( + [[[-2.0, 2.0], [-2.0, 2.0]]] + ) + engine.robot.get_qvel_limits.return_value = torch.full((1, 2), 100.0) + for action in engine.actions.values(): + monkeypatch.setattr( + action, + "_plan", + Mock( + side_effect=AssertionError("Ordinary planning must be bypassed") + ), + ) + self.engine = engine + return engine + + factory = Factory(effect_monitors={}) + adapter = TaskProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + compiled = adapter.compile(_program()) + selector = registry.create_runtime( + scene_provider=_scene_registry().make_scene_provider(batch_size=1), + robot_profile=profile, + control_dt=_STEP_DT, + ) + bridge = adapter.create_bridge(compiled, nominal_plan_factory=selector) + iterator = bridge.iter_segments() + segment = next(iterator) + dispatched = [] + for command in segment.actions: + value = command.value + assert isinstance(value, torch.Tensor) and value.shape == (1, 2) + factory.engine.robot.get_qpos.return_value = value.clone() + dispatched.append(value.clone()) + assert dispatched + assert max(float(x[0, 1]) for x in dispatched) >= 0.19 + assert segment.validator().tolist() == [True] + assert list(iterator) == [] diff --git a/tests/gym/envs/task_program/test_expert_sources.py b/tests/gym/envs/task_program/test_expert_sources.py new file mode 100644 index 000000000..d45246496 --- /dev/null +++ b/tests/gym/envs/task_program/test_expert_sources.py @@ -0,0 +1,327 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Registered expert admission, capabilities and exact call dispatch.""" + +from __future__ import annotations +from dataclasses import dataclass +import gc +import weakref +import torch +import hashlib +from types import SimpleNamespace +from typing import ClassVar +from unittest.mock import Mock +import numpy as np +import pytest + +from embodichain.data_pipeline.real2sim import ( + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, +) +from embodichain.data_pipeline.real2sim.normalization import Real2SimNormalizationCfg +from embodichain.data_pipeline.real2sim.bundle import write_real2sim_bundle +from embodichain.lab.task_program.integrations.expert_sources import ( + ExpertSourceRegistry, + ExpertSourceSelection, +) +from embodichain.lab.task_program.integrations.real2sim import ( + Real2SimExpertSourceProviderFactory, +) +from embodichain.lab.task_program.integrations.configured import ( + _decode_runtime_services, +) +from embodichain.lab.sim.atomic_actions import MotionPolicy, RecoveryPolicy, PickUp + + +def _declaration(tmp_path, positions=None): + (tmp_path / "robot.urdf").write_text("") + q = np.zeros((12, 2)) if positions is None else positions + source = Real2SimTrajectory( + "fixture", + "raw1", + "pick0", + "robot", + "r1", + ("arm", "finger"), + ("rad", "rad"), + q, + np.arange(12) * 0.01, + "s", + "cal1", + "phase1", + "verified", + "verified", + ( + SourcePhase("approach", 0, 6, "free", ("joint_residual",)), + SourcePhase("close", 6, 9, "contact"), + SourcePhase("lift", 9, 12, "hold"), + ), + ) + path = write_real2sim_bundle( + source, + RobotJointLayout( + "robot", "r1", source.joint_names, source.joint_units, q[0], ("arm",) + ), + Real2SimNormalizationCfg(control_dt=0.01), + tmp_path / "data", + ) + return dict( + call_id="pick", + call_index=0, + provider_id="real2sim_qpos", + revision="1", + target_profile_id="fixture", + target_asset_sha256=hashlib.sha256( + (tmp_path / "robot.urdf").read_bytes() + ).hexdigest(), + bundle_path=str(path), + bundle_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + ) + + +def test_configured_source_uses_stable_ids_and_required_capabilities(tmp_path): + value = _declaration(tmp_path) + decoded = _decode_runtime_services({"expert_source_providers": [value]}) + assert decoded.expert_sources.selections[0].call_id == "pick" + with pytest.raises(ValueError, match="capabilities"): + _decode_runtime_services( + { + "expert_source_providers": [ + {**value, "required_capabilities": ["replanning"]} + ] + } + ) + with pytest.raises(ValueError, match="unregistered"): + _decode_runtime_services( + { + "expert_source_providers": [ + {**value, "provider_id": "arbitrary.module:Factory"} + ] + } + ) + with pytest.raises(ValueError, match="unsupported"): + _decode_runtime_services( + {"expert_source_providers": [{**value, "callable": "arbitrary.module.fn"}]} + ) + + +def test_fixed_replay_policy_rejected_before_any_provider_creation(tmp_path): + decoded = _decode_runtime_services( + {"expert_source_providers": [_declaration(tmp_path)]} + ) + with pytest.raises(ValueError, match="max_replans"): + decoded.expert_sources.validate_policy( + SimpleNamespace( + recovery_policy=RecoveryPolicy(), motion_policy=MotionPolicy() + ) + ) + decoded.expert_sources.validate_policy( + SimpleNamespace( + recovery_policy=RecoveryPolicy(max_replans=0, max_action_retries=0), + motion_policy=MotionPolicy(), + ) + ) + + +def test_dispatch_is_exact_and_call_local_with_clock_preflight(tmp_path): + decoded = _decode_runtime_services( + {"expert_source_providers": [_declaration(tmp_path)]} + ) + selector = decoded.expert_sources.create_runtime( + scene_provider=Mock(), + robot_profile=SimpleNamespace(profile_id="fixture"), + control_dt=0.01, + ) + engine = Mock() + engine.robot.cfg.fpath = str(tmp_path / "robot.urdf") + request = SimpleNamespace( + analysis_call_index=0, call=SimpleNamespace(semantic_id="pick") + ) + first = selector.create_nominal_plan_provider(request, engine=engine) + second = selector.create_nominal_plan_provider(request, engine=engine) + assert first.__self__ is not second.__self__ + assert ( + selector.create_nominal_plan_provider( + SimpleNamespace(analysis_call_index=1), engine=engine + ) + is None + ) + with pytest.raises(ValueError, match="kind"): + selector.create_nominal_plan_provider( + SimpleNamespace( + analysis_call_index=0, call=SimpleNamespace(semantic_id="place") + ), + engine=Mock(), + ) + bad = decoded.expert_sources.create_runtime( + scene_provider=Mock(), + robot_profile=SimpleNamespace(profile_id="fixture"), + control_dt=0.02, + ) + with pytest.raises(ValueError, match="clock"): + bad.create_nominal_plan_provider(request, engine=engine) + + +def test_duplicate_provider_or_call_selection_is_not_silently_replaced(tmp_path): + value = _declaration(tmp_path) + factory = Real2SimExpertSourceProviderFactory( + value["bundle_path"], + value["bundle_sha256"], + "fixture", + value["target_asset_sha256"], + ) + selection = ExpertSourceSelection("pick", "real2sim_qpos", "1") + with pytest.raises(ValueError, match="Duplicate expert"): + ExpertSourceRegistry((factory, factory), (selection,)) + with pytest.raises(ValueError, match="occurrence"): + ExpertSourceRegistry((factory,), (selection, selection)) + + +def test_live_target_asset_mismatch_is_rejected(tmp_path): + value = _declaration(tmp_path) + factory = Real2SimExpertSourceProviderFactory( + value["bundle_path"], + value["bundle_sha256"], + "fixture", + value["target_asset_sha256"], + ) + (tmp_path / "robot.urdf").write_text("") + engine = Mock() + engine.robot.cfg.fpath = str(tmp_path / "robot.urdf") + with pytest.raises(ValueError, match="asset digest"): + factory.create( + call_id="pick", + scene_provider=Mock(), + robot_profile=SimpleNamespace(profile_id="fixture"), + control_dt=0.01, + engine=engine, + ) + + +def _selector_and_engine(tmp_path, positions=None): + decoded = _decode_runtime_services( + {"expert_source_providers": [_declaration(tmp_path, positions)]} + ) + selector = decoded.expert_sources.create_runtime( + scene_provider=Mock(), + robot_profile=SimpleNamespace(profile_id="fixture"), + control_dt=0.01, + ) + engine = Mock() + engine.robot.cfg.fpath = str(tmp_path / "robot.urdf") + engine.robot.joint_names = ["arm", "finger"] + engine.robot.get_qpos_limits.return_value = torch.tensor( + [[[-2.0, 2.0], [-2.0, 2.0]]] + ) + engine.robot.get_qvel_limits.return_value = torch.full((1, 2), 2.0) + engine.skills = {"pick_up": PickUp.descriptor()} + request = SimpleNamespace( + analysis_call_index=0, call=SimpleNamespace(semantic_id="pick") + ) + return selector, engine, request + + +@pytest.mark.parametrize( + "amplitude,rejected", [(0.019, False), (0.02, False), (0.03, True)] +) +def test_recorded_source_checks_each_interval_at_reversals( + tmp_path, amplitude, rejected +): + positions = np.zeros((12, 2)) + positions[2, 0] = amplitude + selector, engine, selected = _selector_and_engine(tmp_path, positions) + callback = selector.create_nominal_plan_provider(selected, engine=engine) + provider = callback.__self__ + # Averaged velocities alone pass even though the .03 rad reversal needs + # 3 rad/s over individual 10 ms intervals, above the 2 rad/s live limit. + assert np.abs(provider.unit.velocities).max() <= 2.0 + request = SimpleNamespace( + skill_id="pick_up", + revision=0, + recovery_policy=RecoveryPolicy(max_replans=0, max_action_retries=0), + binding=SimpleNamespace( + endpoints=(SimpleNamespace(target=SimpleNamespace(joint_ids=(0, 1))),) + ), + ) + context = SimpleNamespace( + batch_size=1, + robot=SimpleNamespace(qpos=torch.zeros((1, 2))), + env_ids=torch.tensor([0]), + require_control_dt=lambda: 0.01, + ) + if rejected: + with pytest.raises(ValueError, match="interval speeds.*velocity limits"): + callback(request, context) + engine.build_plan_from_trajectory.assert_not_called() + else: + assert ( + callback(request, context) is engine.build_plan_from_trajectory.return_value + ) + engine.build_plan_from_trajectory.assert_called_once() + + +def test_completed_source_is_collectable_after_callback_is_released(tmp_path): + selector, engine, request = _selector_and_engine(tmp_path) + callback = selector.create_nominal_plan_provider(request, engine=engine) + provider_ref = weakref.ref(callback.__self__) + assert ( + callback.__self__.__hash__ is None + ) # dataclass providers need not be hashable + del callback + gc.collect() + assert provider_ref() is None + # The selector remains live and can serve another call without old arrays. + again = selector.create_nominal_plan_provider(request, engine=engine) + assert again is not None + + +def test_cached_provider_reuse_is_rejected_even_after_its_callback_is_released( + tmp_path, monkeypatch +): + selector, engine, request = _selector_and_engine(tmp_path) + callback = selector.create_nominal_plan_provider(request, engine=engine) + cached = callback.__self__ + monkeypatch.setattr( + type(selector.registry.factories[0]), "create", lambda self, **kwargs: cached + ) + with pytest.raises(ValueError, match="reused"): + selector.create_nominal_plan_provider(request, engine=engine) + del callback + gc.collect() + with pytest.raises(ValueError, match="reused"): + selector.create_nominal_plan_provider(request, engine=engine) + + +def test_non_weak_referenceable_provider_fails_with_explicit_lifecycle_error( + tmp_path, monkeypatch +): + selector, engine, request = _selector_and_engine(tmp_path) + + class SlottedProvider: + __slots__ = () + + def plan(self, request, context): + raise AssertionError("Unsupported lifecycle must fail before planning") + + monkeypatch.setattr( + type(selector.registry.factories[0]), + "create", + lambda self, **kwargs: SlottedProvider(), + ) + with pytest.raises(TypeError, match="weak references"): + selector.create_nominal_plan_provider(request, engine=engine) diff --git a/tests/lab/task_program/test_semantic_executor.py b/tests/lab/task_program/test_semantic_executor.py index 655c10437..6aa8c9c7d 100644 --- a/tests/lab/task_program/test_semantic_executor.py +++ b/tests/lab/task_program/test_semantic_executor.py @@ -941,6 +941,7 @@ def _system( runtime_runner_cfg: ExecutionRunnerCfg | None = None, install_effect_monitor: bool = True, plan_transform_factory: TaskProgramPlanTransformFactory | None = None, + nominal_plan_factory: object | None = None, ) -> _System: robot = Mock() robot.device = torch.device("cpu") @@ -979,6 +980,7 @@ def _system( clock=clock, runner_cfg=runtime_runner_cfg, plan_transform_factory=plan_transform_factory, + nominal_plan_factory=nominal_plan_factory, ) return _System( runtime, @@ -2578,3 +2580,77 @@ def test_split_grasp_and_place_preserve_verified_held_state_across_workflows() - "pick", "place", ] + + +def test_nominal_factory_bypasses_ordinary_planning_and_preserves_effect_verification( + monkeypatch, +): + requests = [] + + class Factory: + def create_nominal_plan_provider(self, request, *, engine): + requests.append(request) + action = engine.actions[request.invocation.skill_id] + + def provide(resolved, context): + goal = action.require_goal(resolved) + position = torch.full( + (context.batch_size, 1), + goal.target_position, + device=context.robot.qpos.device, + ) + return action.build_command_plan( + resolved, + context, + success=goal.plan_success, + commands=TimedCommandSequence((), context.env_ids), + expected_effects=StateDelta( + articulation_joint_updates={ + ("fixture", "joint"): ArticulationJointState(position) + } + ), + effect_verification=EffectVerificationRequirement( + "semantic_effect" + ), + replannable=False, + scene_dependency_monitor_until={"fixture": 0}, + ) + + return provide + + system = _system( + (EffectMonitorDecision(_mask(True, True), _mask(False, False)),), + nominal_plan_factory=Factory(), + ) + + def forbidden(*args, **kwargs): + raise AssertionError("ordinary action planning must not run") + + monkeypatch.setattr(system.action, "_plan", forbidden) + result = system.runtime.start(_call("external"), workflow_id="external-workflow") + for _ in range(10): + if result.terminal: + break + if result.wait_duration: + system.clock.sleep(result.wait_duration) + result = system.runtime.step() + assert system.action.plan_count == 0 + assert len(requests) == 1 and requests[0].workflow_id == "external-workflow" + assert requests[0].workflow_call_index == 0 + assert result.success_mask.tolist() == [True, True], result.to_metadata() + assert result.effects + + +def test_nominal_factory_rejects_noncallable_provider(): + class Factory: + def create_nominal_plan_provider(self, request, *, engine): + return object() + + system = _system( + (EffectMonitorDecision(_mask(True, True), _mask(False, False)),), + nominal_plan_factory=Factory(), + ) + result = system.runtime.start(_call("external"), workflow_id="external-workflow") + assert result.terminal and not result.success_mask.any() + assert "callable or None" in json.dumps(result.to_metadata()) + assert system.action.plan_count == 0 diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index eb92a08c9..60729359f 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -1849,7 +1849,7 @@ def test_planner_timing_is_preserved_in_simple_action() -> None: action = _bind_action(generator, MoveJoints()) invocation = ActionInvocation( skill_id="move_joints", - goal=JointPositionGoal(torch.ones(ARM_DOF)), + goal=JointPositionGoal(torch.full((ARM_DOF,), 0.15)), binding=_binding(action), motion_policy=MotionPolicy(strategy="motion_gen", sample_count=3), ) @@ -7367,3 +7367,306 @@ def test_assemble_place_reports_sampled_symmetry() -> None: assert plan.diagnostics.metadata["affordance_sample"]["assembly"][ "candidate_ids" ] == [0, 1] + + +def test_recorded_pick_builds_commands_tracking_effects_and_preserves_rebuild( + monkeypatch, +): + from embodichain.lab.sim.atomic_actions.plans import PlannerDiagnostics + from embodichain.lab.sim.atomic_actions.policies import RecoveryPolicy + + generator = _motion_generator() + generator.robot.get_qpos_limits.return_value = torch.tensor([-2.0, 2.0]).expand( + NUM_ENVS, ROBOT_DOF, 2 + ) + generator.robot.get_qvel_limits.return_value = torch.full( + (NUM_ENVS, ROBOT_DOF), 100.0 + ) + action = _bind_action(generator, PickUp()) + engine = _ACTION_ENGINES[id(action)] + monkeypatch.setattr( + action, + "_plan", + Mock(side_effect=AssertionError("ordinary planning must be bypassed")), + ) + monkeypatch.setattr( + generator.robot, + "compute_fk", + lambda qpos, **kwargs: torch.eye(4).repeat(qpos.shape[0], 1, 1), + ) + context = _context( + scene=_target_scene( + torch.eye(4).repeat(NUM_ENVS, 1, 1), timestamp=0.0, version=0 + ) + ) + invocation = ActionInvocation( + skill_id="pick_up", + goal=GraspGoal(semantics=_semantics(entity_id="target")), + binding=_binding(action), + recovery_policy=RecoveryPolicy(max_replans=0, max_action_retries=0), + ) + request = engine.resolve(invocation) + q = torch.zeros(NUM_ENVS, 12, ROBOT_DOF) + q[:, 1:, 0] = torch.linspace(0.01, 0.11, 11) + q[:, 6:, ARM_DOF:] = 0.2 + trajectory = TimedTrajectory.from_uniform_step( + q, env_ids=context.env_ids, step_dt=CONTROL_DT + ) + plan = engine.plan_request( + request, + context, + nominal_plan_provider=lambda req, ctx: engine.build_plan_from_recorded_trajectory( + req, + ctx, + trajectory, + segment_lengths={"approach": 6, "close": 3, "lift": 3}, + diagnostics=PlannerDiagnostics( + "real2sim_qpos", metadata={"source_revision": "fixture1"} + ), + ), + ) + assert [s.name for s in plan.segments] == ["approach", "close", "lift"] + assert plan.commands.frame_count == 12 and plan.tracking is not None + assert not plan.replannable + assert context.get_held_object("arm") is None + assert ( + plan.expected_effects.apply(context.task, plan.plan_success).get_held_object( + "arm" + ) + is not None + ) + torch.testing.assert_close( + _joint_command_positions(plan, "hand")[:, 6:], q[:, 6:, ARM_DOF:] + ) + assert plan.scene_dependency_monitor_until["target"] == 6 + rebuilt = engine.rebuild_plan_from_trajectory(request, context, plan, trajectory) + assert [s.name for s in rebuilt.segments] == ["approach", "close", "lift"] + assert rebuilt.diagnostics.metadata["source_revision"] == "fixture1" + assert not rebuilt.replannable + assert rebuilt.trajectory_input.geometry is None + + def shifted_fk(qpos: torch.Tensor, **_: object) -> torch.Tensor: + pose = torch.eye(4).repeat(qpos.shape[0], 1, 1) + pose[:, 0, 3] = qpos[:, 0] + return pose + + monkeypatch.setattr(generator.robot, "compute_fk", shifted_fk) + changed = trajectory.positions.clone() + grasp_index = plan.trajectory_input.landmarks["grasp"] + changed[:, grasp_index, 0] += 0.001 + candidate = replace(trajectory, positions=changed, velocities=None) + rebuilt = engine.rebuild_plan_from_trajectory(request, context, plan, candidate) + torch.testing.assert_close( + rebuilt.expected_effects.held_object_updates["arm"].object_to_eef[:, 0, 3], + changed[:, grasp_index, 0], + ) + + +@pytest.mark.parametrize("skill", ["pick_up", "move_joints"]) +def test_normal_and_external_motion_share_semantic_assembly( + skill: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from embodichain.lab.sim.atomic_actions import RecoveryPolicy + from embodichain.lab.task_program.integrations.catalog import _canonical_value + + generator = _motion_generator() + generator.robot.get_qpos_limits.return_value = torch.tensor([-100.0, 100.0]).expand( + NUM_ENVS, ROBOT_DOF, 2 + ) + generator.robot.get_qvel_limits.return_value = torch.full( + (NUM_ENVS, ROBOT_DOF), 10000.0 + ) + action = _bind_action(generator, PickUp() if skill == "pick_up" else MoveJoints()) + engine = _ACTION_ENGINES[id(action)] + goal = ( + GraspGoal( + _semantics(entity_id="target"), + grasp_xpos=torch.eye(4).repeat(NUM_ENVS, 1, 1), + ) + if skill == "pick_up" + else JointPositionGoal(torch.full((ARM_DOF,), 0.1)) + ) + context = _context( + scene=_target_scene( + torch.eye(4).repeat(NUM_ENVS, 1, 1), timestamp=0.0, version=0 + ) + ) + invocation = replace( + _invocation(action, goal), + recovery_policy=RecoveryPolicy(max_replans=0, max_action_retries=0), + ) + request = engine.resolve(invocation) + normal = engine.plan_request(request, context) + assert normal.trajectory_input is not None + monkeypatch.setattr( + action, + "_plan", + Mock(side_effect=AssertionError("ordinary planning must be bypassed")), + ) + imported = engine.build_plan_from_trajectory( + request, context, replace(normal.trajectory_input, replannable=False) + ) + assert _canonical_value(imported.commands) == _canonical_value(normal.commands) + assert _canonical_value(imported.tracking) == _canonical_value(normal.tracking) + assert imported.segments == normal.segments + assert ( + imported.scene_dependency_monitor_until == normal.scene_dependency_monitor_until + ) + assert not imported.replannable and normal.replannable + assert context.get_held_object("arm") is None + if skill == "pick_up": + expected = normal.expected_effects.held_object_updates["arm"] + actual = imported.expected_effects.held_object_updates["arm"] + torch.testing.assert_close(actual.object_to_eef, expected.object_to_eef) + torch.testing.assert_close(actual.grasp_xpos, expected.grasp_xpos) + positions = normal.joint_trajectory.positions.clone() + positions[:, normal.trajectory_input.landmarks["grasp"], 0] += 0.001 + with pytest.raises(ValueError, match="new trajectory assembly input"): + engine.rebuild_plan_from_trajectory( + request, + context, + normal, + replace(normal.joint_trajectory, positions=positions), + ) + else: + assert imported.expected_effects.is_empty + + +def test_real2sim_provider_constructs_one_pick_and_rejects_wrong_live_initial_state( + tmp_path, monkeypatch +): + import hashlib + import numpy as np + from embodichain.data_pipeline.real2sim import ( + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, + ) + from embodichain.data_pipeline.real2sim.bundle import write_real2sim_bundle + from embodichain.data_pipeline.real2sim.normalization import ( + Real2SimNormalizationCfg, + ) + from embodichain.lab.task_program.integrations.real2sim import ( + Real2SimExpertSourceProviderFactory, + ) + from embodichain.lab.sim.atomic_actions import RecoveryPolicy + + asset = tmp_path / "robot.urdf" + asset.write_text("") + generator = _motion_generator() + generator.robot.cfg.fpath = str(asset) + generator.robot.num_instances = 1 + generator.robot.get_qpos.side_effect = lambda name=None: torch.zeros( + 1, ARM_DOF if name == "arm" else HAND_DOF if name == "hand" else ROBOT_DOF + ) + action = _bind_action(generator, PickUp()) + engine = _ACTION_ENGINES[id(action)] + monkeypatch.setattr( + generator.robot, + "compute_fk", + lambda qpos, **kwargs: torch.eye(4).repeat(qpos.shape[0], 1, 1), + ) + monkeypatch.setattr( + generator.robot, + "get_qpos_limits", + lambda **kwargs: torch.tensor([[[-2.0, 2.0]] * ROBOT_DOF]), + ) + monkeypatch.setattr( + generator.robot, + "get_qvel_limits", + lambda **kwargs: torch.full((1, ROBOT_DOF), 100.0), + ) + monkeypatch.setattr( + action, "_plan", Mock(side_effect=AssertionError("No ordinary planning")) + ) + q = np.zeros((12, ROBOT_DOF)) + q[1:, 0] = np.linspace(0.01, 0.1, 11) + q[6:, ARM_DOF:] = 0.2 + raw = Real2SimTrajectory( + "synthetic", + "source1", + "pick0", + "fixture_robot", + "model1", + tuple(generator.robot.joint_names), + ("rad",) * ROBOT_DOF, + q, + np.arange(12) * CONTROL_DT, + "s", + "cal1", + "ann1", + "verified", + "verified", + ( + SourcePhase("approach", 0, 6, "free", ("joint_residual",)), + SourcePhase("close", 6, 9, "contact"), + SourcePhase("lift", 9, 12, "hold"), + ), + ) + path = write_real2sim_bundle( + raw, + RobotJointLayout( + raw.robot_id, + raw.robot_revision, + raw.joint_names, + raw.joint_units, + q[0], + raw.joint_names[:ARM_DOF], + ), + Real2SimNormalizationCfg(control_dt=CONTROL_DT), + tmp_path / "source", + ) + factory = Real2SimExpertSourceProviderFactory( + str(path), + hashlib.sha256(path.read_bytes()).hexdigest(), + "profile1", + hashlib.sha256(asset.read_bytes()).hexdigest(), + ) + provider = factory.create( + call_id="pick", + scene_provider=Mock(), + robot_profile=SimpleNamespace(profile_id="profile1"), + control_dt=CONTROL_DT, + engine=engine, + ) + ctx = PlanningContext( + robot=RobotObservation( + timestamp=0, qpos=torch.zeros(1, ROBOT_DOF), qvel=torch.zeros(1, ROBOT_DOF) + ), + task=TaskState.empty(1, "cpu"), + scene=_target_scene(torch.eye(4)[None], timestamp=0, version=0), + env_ids=torch.tensor([0]), + control_dt=CONTROL_DT, + ) + req = engine.resolve( + ActionInvocation( + skill_id="pick_up", + goal=GraspGoal(semantics=_semantics(entity_id="target")), + binding=_binding(action), + recovery_policy=RecoveryPolicy(max_replans=0, max_action_retries=0), + ) + ) + plan = engine.plan_request(req, ctx, nominal_plan_provider=provider.plan) + assert plan.commands.frame_count == 12 and plan.plan_success.tolist() == [True] + assert ( + plan.expected_effects.apply(ctx.task, plan.plan_success).get_held_object("arm") + is not None + ) + assert plan.diagnostics.metadata["source_revision"] == "source1" + assert not plan.replannable + wrong = replace( + ctx, + robot=RobotObservation( + timestamp=0, qpos=torch.ones(1, ROBOT_DOF), qvel=torch.zeros(1, ROBOT_DOF) + ), + ) + with pytest.raises(ValueError, match="initial state"): + provider.plan(req, wrong) + monkeypatch.setattr( + generator.robot, + "get_qvel_limits", + lambda **kwargs: torch.full((1, ROBOT_DOF), 0.001), + ) + with pytest.raises(ValueError, match="velocity limits"): + provider.plan(req, ctx) diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index c59806540..b3ad862c3 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -928,3 +928,54 @@ def test_rebuild_plan_validates_source_plan_identity() -> None: replace(plan, invocation_id="different"), plan.joint_trajectory, ) + + +def test_nominal_provider_bypasses_ordinary_planning_and_owns_result(monkeypatch): + engine = _engine() + action = StubAction() + engine.register(action) + invocation = _invocation(engine, torch.ones(2, 3)) + retained = [] + + def provide(request, context): + trajectory = TimedTrajectory.from_uniform_step( + torch.stack([context.robot.qpos, context.robot.qpos + 0.1], dim=1), + env_ids=context.env_ids, + step_dt=ACTION_DT, + ) + result = action.build_plan( + request, context, success=True, trajectory=trajectory, replannable=False + ) + retained.append(result) + return result + + monkeypatch.setattr( + action, + "_plan", + Mock(side_effect=AssertionError("ordinary planning must not run")), + ) + result = engine.plan(invocation, nominal_plan_provider=provide) + assert result.success_all and not result.replannable + retained[0].joint_trajectory.positions.fill_(42.0) + assert result.joint_trajectory.positions[0, -1, 0] == pytest.approx(0.1) + action._plan.assert_not_called() + + +def test_nominal_provider_failure_and_invalid_result_do_not_fall_back(monkeypatch): + engine = _engine() + action = StubAction() + engine.register(action) + monkeypatch.setattr( + action, "_plan", Mock(side_effect=AssertionError("unexpected fallback")) + ) + invocation = _invocation(engine, torch.ones(2, 3)) + with pytest.raises(RuntimeError, match="source failure"): + engine.plan( + invocation, + nominal_plan_provider=Mock(side_effect=RuntimeError("source failure")), + ) + with pytest.raises(TypeError, match="return an ActionPlan"): + engine.plan(invocation, nominal_plan_provider=lambda request, context: None) + with pytest.raises(TypeError, match="callable"): + engine.plan(invocation, nominal_plan_provider=object()) + action._plan.assert_not_called() diff --git a/tests/sim/atomic_actions/test_engine_per_env.py b/tests/sim/atomic_actions/test_engine_per_env.py index b2543bc65..503233599 100644 --- a/tests/sim/atomic_actions/test_engine_per_env.py +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -3525,3 +3525,38 @@ def test_failed_effect_plan_retries_without_requesting_effect_verification() -> assert any( event.kind is ExecutionEventKind.RECOVERY_EXHAUSTED for event in failed.events ) + + +def test_nominal_source_is_retained_for_recovery_and_invocation_revision(monkeypatch): + engine, action = _engine() + calls = [] + + def provide(request, context): + calls.append((request.revision, context.scene.version)) + target = context.robot.qpos + 0.1 + trajectory = TimedTrajectory.from_uniform_step( + torch.stack( + [ + torch.lerp(context.robot.qpos, target, t) + for t in torch.linspace(0, 1, 20) + ], + dim=1, + ), + env_ids=context.env_ids, + step_dt=CONTROL_DT, + ) + return action.build_plan(request, context, success=True, trajectory=trajectory) + + def forbidden(*args, **kwargs): + raise AssertionError("ordinary planning must not run") + + monkeypatch.setattr(action, "_plan", forbidden) + initial = _context(0.0, 0.0, 0.1, 0) + invocation = _invocation(engine) + session = engine.start((invocation,), initial, nominal_plan_provider=provide) + session.tick(initial) + changed = _context(0.1, 0.0, 0.3, 1) + session.tick(changed) + session.revise_current(replace(invocation, revision=1), context=changed) + assert calls == [(0, 0), (0, 1), (1, 1)] + assert action.plan_count == 0 diff --git a/tests/sim/atomic_actions/test_trajectory_assembly.py b/tests/sim/atomic_actions/test_trajectory_assembly.py new file mode 100644 index 000000000..7c3d8d41a --- /dev/null +++ b/tests/sim/atomic_actions/test_trajectory_assembly.py @@ -0,0 +1,370 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Source-neutral trajectory schemas, snapshots and complete Skill assembly.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import ClassVar +from unittest.mock import Mock + +import pytest +import torch + +from embodichain.lab.sim.atomic_actions import ( + ActionInvocation, + ActionPlan, + AtomicAction, + AtomicActionEngine, + JointPositionGoal, + MoveJoints, + Place, + PlannerDiagnostics, + RecoveryPolicy, + SkillPlanSemantics, + SkillTrajectoryContract, + SkillTrajectoryInput, + TimedTrajectory, + TrajectorySegment, +) +from embodichain.lab.sim.atomic_actions.primitives.pick_up import PickTrajectoryGeometry +from embodichain.lab.task_program.integrations.catalog import _canonical_value + +CONTROL_DT = 0.01 +BATCH_SIZE = 2 +ROBOT_DOF = 3 +SAMPLE_COUNT = 5 + + +def _motion(**changes: object) -> SkillTrajectoryInput: + positions = torch.zeros(BATCH_SIZE, SAMPLE_COUNT, ROBOT_DOF) + positions[:, :, :2] = torch.linspace(0, 0.01, SAMPLE_COUNT)[None, :, None] + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.tensor([7, 3]), + step_dt=CONTROL_DT, + ) + values = dict( + schema_id="move_joints.qpos", + schema_revision="1", + trajectory=trajectory, + landmarks={"goal": SAMPLE_COUNT - 1}, + diagnostics=PlannerDiagnostics("external_fixture"), + ) + values.update(changes) + return SkillTrajectoryInput(**values) + + +def _engine() -> tuple[AtomicActionEngine, object, object]: + robot = Mock() + robot.dof = ROBOT_DOF + robot.device = torch.device("cpu") + robot.joint_names = ["shoulder", "elbow", "finger"] + robot.mimic_ids = [] + robot.control_parts = {"custom_motion": object(), "custom_grasp": object()} + robot.get_qpos.return_value = torch.zeros(BATCH_SIZE, ROBOT_DOF) + robot.get_qvel.return_value = torch.zeros(BATCH_SIZE, ROBOT_DOF) + robot.get_joint_ids.side_effect = lambda name=None, **_: ( + [0, 1] + if name == "custom_motion" + else [2] if name == "custom_grasp" else [0, 1, 2] + ) + robot.get_qpos_limits.return_value = torch.tensor([-2.0, 2.0]).expand( + BATCH_SIZE, ROBOT_DOF, 2 + ) + robot.get_qvel_limits.return_value = torch.full((BATCH_SIZE, ROBOT_DOF), 2.0) + generator = Mock(robot=robot, device=torch.device("cpu")) + generator.planner.cfg.planner_type = "fixture" + generator.planner.collision_world_info = None + engine = AtomicActionEngine(generator, load_builtins=False) + engine.register(MoveJoints()) + invocation = ActionInvocation( + skill_id="move_joints", + goal=JointPositionGoal(torch.full((BATCH_SIZE, 2), 0.01)), + binding=engine.bind_control_parts( + "move_joints", {"primary": {"motion": "custom_motion"}} + ), + recovery_policy=RecoveryPolicy(max_replans=0, max_action_retries=0), + ) + request = engine.resolve(invocation) + context = replace( + engine.initial_context(), env_ids=torch.tensor([7, 3]), control_dt=CONTROL_DT + ) + return engine, request, context + + +def test_assembly_input_and_plan_snapshots_own_all_mutable_values() -> None: + engine, request, context = _engine() + motion = _motion() + plan = engine.build_plan_from_trajectory(request, context, motion) + copied = plan.snapshot() + motion.trajectory.positions.fill_(0.5) + motion.planning_success.zero_() + plan.trajectory_input.trajectory.positions.fill_(0.7) + assert copied.joint_trajectory.positions[:, -1, 0].tolist() == pytest.approx( + [0.01, 0.01] + ) + assert copied.trajectory_input.trajectory.positions[ + :, -1, 0 + ].tolist() == pytest.approx([0.01, 0.01]) + assert copied.plan_success.all() + with pytest.raises(TypeError): + copied.trajectory_input.landmarks["goal"] = 0 + + +def test_geometry_is_typed_and_owned_without_becoming_measured_evidence() -> None: + pose = torch.eye(4)[None] + geometry = PickTrajectoryGeometry(pose) + motion = _motion(geometry=geometry) + pose[:, 0, 3] = 1 + geometry.grasp_pose[:, 0, 3] = 2 + assert motion.geometry.grasp_pose[0, 0, 3] == 0 + with pytest.raises(TypeError, match="geometric facts"): + MoveJoints.trajectory_contract.validate_input(motion) + + +@pytest.mark.parametrize( + "changes,error", + [ + ({"schema_revision": "2"}, "schema identity or revision"), + ({"schema_id": "pick_up.qpos"}, "schema identity or revision"), + ({"landmarks": {}}, "event landmarks"), + ({"landmarks": {"goal": 0}}, "stage boundary"), + ( + {"segments": (TrajectorySegment("other", 0, SAMPLE_COUNT),)}, + "semantic stages", + ), + ], +) +def test_skill_contract_rejects_incompatible_input(changes: dict, error: str) -> None: + with pytest.raises(ValueError, match=error): + MoveJoints.trajectory_contract.validate_input(_motion(**changes)) + + +@pytest.mark.parametrize( + "landmarks", [{"goal": -1}, {"goal": SAMPLE_COUNT}, {"goal": True}] +) +def test_event_indices_are_actual_samples(landmarks: dict) -> None: + with pytest.raises(ValueError, match="actual trajectory samples"): + _motion(landmarks=landmarks) + + +@pytest.mark.parametrize( + "segments", + [ + (TrajectorySegment("a", 1, SAMPLE_COUNT),), + (TrajectorySegment("a", 0, 2), TrajectorySegment("b", 3, SAMPLE_COUNT)), + (TrajectorySegment("a", 0, 3), TrajectorySegment("b", 2, SAMPLE_COUNT)), + (TrajectorySegment("a", 0, SAMPLE_COUNT - 1),), + (TrajectorySegment("a", 0, 2), TrajectorySegment("a", 2, SAMPLE_COUNT)), + ], +) +def test_semantic_stages_cover_path_without_gaps_or_repeated_ids( + segments: tuple, +) -> None: + with pytest.raises(ValueError): + _motion(segments=segments) + + +def test_declared_stage_events_follow_current_grid_and_are_fingerprinted() -> None: + contract = SkillTrajectoryContract( + "two_stage", + "1", + ("approach", "close"), + ("grasp",), + {"grasp": "close"}, + ) + short = (TrajectorySegment("approach", 0, 2), TrajectorySegment("close", 2, 5)) + long = (TrajectorySegment("approach", 0, 4), TrajectorySegment("close", 4, 9)) + assert contract.stage_landmarks(short, 5)["grasp"] == 4 + assert contract.stage_landmarks(long, 9)["grasp"] == 8 + descriptor = MoveJoints.descriptor() + assert _canonical_value(descriptor) != _canonical_value( + replace( + descriptor, + trajectory_contract=replace(descriptor.trajectory_contract, revision="2"), + ) + ) + + +def test_complete_external_move_plan_preserves_custom_binding_and_env_order() -> None: + engine, request, context = _engine() + engine.actions["move_joints"]._plan = Mock( + side_effect=AssertionError("ordinary planning must not run") + ) + plan = engine.build_plan_from_trajectory(request, context, _motion()) + assert plan.commands.frame_count == SAMPLE_COUNT + assert plan.tracking is not None + assert [target.target_id for target in plan.commands.targets] == ["custom_motion"] + assert plan.commands.env_ids.tolist() == [7, 3] + assert [stage.name for stage in plan.segments] == ["move_joints"] + assert plan.expected_effects.is_empty and not plan.replannable + assert plan.trajectory_input.landmarks["goal"] == SAMPLE_COUNT - 1 + + +@pytest.mark.parametrize( + "problem,error", + [ + ("initial", "initial state"), + ("position", "position limits"), + ("velocity", "velocity limits"), + ("reversal", "interval speeds"), + ("endpoint", "authorized endpoints"), + ("goal", "requested joint goal"), + ("env_order", "ordered env_ids"), + ("clock", "control grid"), + ("long_clock", "control grid"), + ("zero_interval", "positive after"), + ], +) +def test_external_motion_admission_fails_before_execution( + problem: str, error: str +) -> None: + engine, request, context = _engine() + motion = _motion() + trajectory = motion.trajectory + q = trajectory.positions.clone() + if problem == "initial": + q[:, 0, 0] = 0.005 + elif problem == "position": + q[:, 2, 0] = 3.0 + elif problem == "velocity": + trajectory = replace(trajectory, velocities=torch.full_like(q, 3.0)) + elif problem == "reversal": + q[:, 2, 0] = 0.03 + trajectory = replace(trajectory, velocities=torch.zeros_like(q)) + elif problem == "endpoint": + q[:, 1:, 2] = 0.005 + elif problem == "goal": + q[:, -1, 0] = 0.008 + elif problem == "env_order": + trajectory = replace(trajectory, env_ids=trajectory.env_ids.flip(0)) + elif problem == "clock": + trajectory = replace(trajectory, dt=trajectory.dt * 1.5) + elif problem == "long_clock": + dt = trajectory.dt.clone() + dt[:, 1] = 1000.0025 # A quarter tick must remain invalid for a long hold. + trajectory = replace(trajectory, dt=dt) + elif problem == "zero_interval": + dt = trajectory.dt.clone() + dt[:, 1] = 0 + trajectory = replace(trajectory, dt=dt, velocities=torch.zeros_like(q)) + motion = replace(motion, trajectory=replace(trajectory, positions=q)) + with pytest.raises(ValueError, match=error): + engine.build_plan_from_trajectory(request, context, motion) + + +def test_move_checks_ordered_waypoints_and_failed_rows_remain_held() -> None: + engine, request, context = _engine() + goals = torch.tensor([[[0.0075, 0.0075], [0.0025, 0.0025], [0.01, 0.01]]]).expand( + BATCH_SIZE, -1, -1 + ) + request = replace(request, goal=JointPositionGoal(goals)) + with pytest.raises(ValueError, match="waypoints in order"): + engine.build_plan_from_trajectory(request, context, _motion()) + request = replace( + request, goal=JointPositionGoal(torch.full((BATCH_SIZE, 2), 0.01)) + ) + motion = _motion(planning_success=torch.tensor([True, False])) + plan = engine.build_plan_from_trajectory(request, context, motion) + assert plan.plan_success.tolist() == [True, False] + assert torch.count_nonzero(plan.joint_trajectory.positions[1]) == 0 + assert plan.trajectory_input.planning_success.tolist() == [True, False] + + +def test_move_accepts_ordered_polyline_waypoints_between_control_samples() -> None: + engine, request, context = _engine() + goals = torch.tensor([[[0.003, 0.003], [0.006, 0.006], [0.01, 0.01]]]).expand( + BATCH_SIZE, -1, -1 + ) + plan = engine.build_plan_from_trajectory( + replace(request, goal=JointPositionGoal(goals)), context, _motion() + ) + assert plan.plan_success.all() + + +def test_same_grid_variant_is_reassembled_and_changed_grid_is_rejected() -> None: + engine, request, context = _engine() + motion = _motion() + plan = engine.build_plan_from_trajectory(request, context, motion) + plan = replace( + plan, + diagnostics=PlannerDiagnostics( + "external_fixture", metadata={"candidate_lineage": "seed_7"} + ), + ) + positions = plan.joint_trajectory.positions.clone() + positions[:, 2, 0] += 0.001 + candidate = replace(plan.joint_trajectory, positions=positions, velocities=None) + rebuilt = engine.rebuild_plan_from_trajectory(request, context, plan, candidate) + assert rebuilt.trajectory_input.landmarks == motion.landmarks + assert rebuilt.diagnostics.backend == "external_fixture" + assert rebuilt.diagnostics.metadata["candidate_lineage"] == "seed_7" + assert rebuilt.commands.frame_count == SAMPLE_COUNT + assert rebuilt.commands.frames[2].commands[0].payload.positions[ + :, 0 + ].tolist() == pytest.approx([0.006, 0.006]) + with pytest.raises(ValueError, match="identical dt"): + engine.rebuild_plan_from_trajectory( + request, context, plan, replace(candidate, dt=candidate.dt * 2) + ) + + +def test_contract_can_limit_batch_without_limiting_other_skills() -> None: + contract = replace(MoveJoints.trajectory_contract, supports_batch=False) + with pytest.raises(ValueError, match="B=1"): + contract.validate_input(_motion()) + + +def test_unsupported_skill_rejects_motion_without_planner_fallback() -> None: + engine, request, context = _engine() + action = Place() + action._plan = Mock(side_effect=AssertionError("ordinary planning must not run")) + engine.register(action) + with pytest.raises(ValueError, match="does not support trajectory assembly"): + engine.build_plan_from_trajectory( + replace(request, skill_id="place"), context, _motion() + ) + action._plan.assert_not_called() + + +def test_unimplemented_contract_and_framework_override_are_rejected() -> None: + class MissingAssembler(AtomicAction): + skill_id: ClassVar[str] = "missing_assembler" + GoalType: ClassVar[type] = JointPositionGoal + trajectory_contract: ClassVar[SkillTrajectoryContract] = ( + SkillTrajectoryContract("missing", "1") + ) + + def _plan(self, request, context) -> ActionPlan: + raise AssertionError("planning must not run") + + engine, _, _ = _engine() + with pytest.raises(ValueError, match="semantic assembler"): + engine.register(MissingAssembler()) + with pytest.raises(TypeError, match="framework-owned"): + + class OverrideAssembler(MoveJoints): + def assemble_trajectory_plan(self, request, context, motion) -> ActionPlan: + raise AssertionError + + +def test_semantic_result_owns_scene_bounds() -> None: + bounds = {"cube": 3} + semantics = SkillPlanSemantics(scene_dependency_monitor_until=bounds) + bounds["cube"] = 4 + assert semantics.scene_dependency_monitor_until["cube"] == 3 diff --git a/tests/sim/motion/expansion/test_real2sim.py b/tests/sim/motion/expansion/test_real2sim.py new file mode 100644 index 000000000..0a0abd37c --- /dev/null +++ b/tests/sim/motion/expansion/test_real2sim.py @@ -0,0 +1,212 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Real2sim templates use the public candidate contracts and local RNG only.""" + +from __future__ import annotations + +from dataclasses import replace +import numpy as np +import pytest +import torch + +from embodichain.data_pipeline.real2sim import ( + Real2SimTrajectory, + RobotJointLayout, + SourcePhase, + MimicJoint, +) +from embodichain.data_pipeline.real2sim.normalization import ( + Real2SimNormalizationCfg, + normalize_real2sim, +) +from embodichain.compute.trajectory import differentiate_positions +from embodichain.lab.sim.motion.expansion import ( + Real2SimSourceAdapter, + SourceContext, + SourceAdapter, + SceneCase, +) +from embodichain.lab.sim.motion.expansion.operators import joint_residual + + +def _unit(): + q = np.c_[ + np.linspace(0.0, 0.3, 20), + np.linspace(0.0, 0.1, 20) ** 2, + np.r_[np.zeros(10), np.linspace(0.0, 0.5, 5), np.full(5, 0.5)], + ] + raw = Real2SimTrajectory( + "fixture", + "r1", + "pick-0", + "robot", + "v1", + ("arm", "elbow", "finger"), + ("rad", "rad", "rad"), + q, + np.arange(20) * 0.01, + "s", + "cal1", + "ph1", + "verified", + "verified", + ( + SourcePhase("approach", 0, 10, "free", ("joint_residual",)), + SourcePhase("close", 10, 15, "contact"), + SourcePhase("hold", 15, 20, "hold"), + ), + ) + layout = RobotJointLayout( + "robot", + "v1", + ("arm", "elbow", "finger", "mimic"), + ("rad", "rad", "rad", "rad"), + np.zeros(4), + ("arm", "elbow"), + (MimicJoint("mimic", "finger", -1.0),), + ) + return normalize_real2sim(raw, layout, Real2SimNormalizationCfg(control_dt=0.01)) + + +def _context(unit): + return SourceContext( + unit.source.source_id, + unit.revision, + unit.source.unit_id, + SceneCase("case", "initial", "scene1", "pick", "profile1"), + 0.01, + ) + + +def test_two_variants_preserve_contacts_hand_mimics_and_boundary_velocities(): + unit = _unit() + adapter = Real2SimSourceAdapter() + assert isinstance(adapter, SourceAdapter) + template = adapter.export_template(unit, context=_context(unit)) + limits = torch.tensor( + [[-2.0, 2.0], [-2.0, 2.0], [0.0, 1.0], [-1.0, 0.0]], + dtype=template.positions.dtype, + ) + variants = [] + for seed in (7, 11): + variant = joint_residual( + template, + joint_limits=limits, + normalized_scale=0.002, + generator=torch.Generator().manual_seed(seed), + ) + repeated = joint_residual( + template, + joint_limits=limits, + normalized_scale=0.002, + generator=torch.Generator().manual_seed(seed), + ) + assert torch.equal(variant.positions, repeated.positions) + assert torch.equal(variant.positions[10:], template.positions[10:]) + assert torch.equal(variant.positions[:, 2:], template.positions[:, 2:]) + assert torch.equal( + variant.positions[[0, 1, 8, 9]], template.positions[[0, 1, 8, 9]] + ) + v = differentiate_positions(variant.positions[None], variant.dt[None])[0] + original = differentiate_positions(template.positions[None], template.dt[None])[ + 0 + ] + assert torch.equal(v[[0, 9, 10, 14, 15, 19]], original[[0, 9, 10, 14, 15, 19]]) + variants.append(variant) + assert not torch.equal(variants[0].positions, variants[1].positions) + np.testing.assert_array_equal(unit.positions, template.positions.numpy()) + assert unit.phases[0].start == 0 and unit.phases[0].stop == 10 + assert template.phases[0].start_index == 1 and template.phases[0].stop_index == 9 + + +def test_unverified_clock_mismatch_and_unsupported_retime_fail_closed(): + unit = _unit() + context = _context(unit) + adapter = Real2SimSourceAdapter() + with pytest.raises(ValueError, match="control clocks"): + adapter.export_template(unit, context=replace(context, control_dt=0.02)) + with pytest.raises(ValueError, match="identity"): + adapter.export_template(unit, context=replace(context, source_revision="wrong")) + unverified = replace( + unit, report={**unit.report_dict(), "input_prerequisites_verified": False} + ) + with pytest.raises(ValueError, match="Unverified"): + adapter.export_template(unverified, context=context) + phase = replace(unit.phases[0], allowed_operators=("retime",)) + with pytest.raises(ValueError, match="same-grid"): + adapter.export_template( + replace(unit, phases=(phase, *unit.phases[1:])), context=context + ) + + +def test_coupled_driver_expansion_is_not_silently_inconsistent(): + unit = _unit() + with pytest.raises(ValueError, match="mimic-driver"): + Real2SimSourceAdapter().export_template( + replace(unit, expandable_joint_indices=(2,)), context=_context(unit) + ) + + +def test_existing_coordinator_generates_nominal_and_two_deterministic_variants(): + from embodichain.lab.sim.motion.expansion import ( + CandidateCoordinator, + ExpansionSession, + TrajectoryExpansionJobCfg, + ) + + unit = _unit() + context = _context(unit) + cfg = TrajectoryExpansionJobCfg.from_mapping( + { + "augmentation": { + "seed": 7, + "max_variants_per_reference": 3, + "factors": { + "spatial": { + "enabled": True, + "method": ["joint_residual"], + "joint_offset_scale": 0.002, + } + }, + "coverage": {"joint_dedup_normalized_tol": 1e-5}, + }, + "scheduling": {"candidate_budget": 8}, + } + ) + limits = torch.tensor( + [[-2.0, 2.0], [-2.0, 2.0], [0.0, 1.0], [-1.0, 0.0]], dtype=torch.float64 + ) + + def generate(): + session = ExpansionSession(cfg) + session.register_case(context.scene_case, limits, joint_names=unit.joint_names) + coordinator = CandidateCoordinator( + cfg, + session=session, + source_adapter=Real2SimSourceAdapter(), + source_context=context, + joint_limits=limits, + ) + return coordinator.enqueue_source(unit, count=3), session + + first, session = generate() + second, _ = generate() + assert len(first) == len(second) == 3 + for left, right in zip(first, second): + assert left.spec.identity == right.spec.identity + assert torch.equal(left.template.positions, right.template.positions) + assert session.snapshot()["counts"]["proposed"] == 3 diff --git a/tests/sim/utility/test_action_utils.py b/tests/sim/utility/test_action_utils.py index fa27a2556..d7c0a43fa 100644 --- a/tests/sim/utility/test_action_utils.py +++ b/tests/sim/utility/test_action_utils.py @@ -21,7 +21,20 @@ from embodichain.compute import trajectory from embodichain.lab.sim.utility import action_utils +# The compatibility surface preserves historical names. New compute APIs are +# consumed through compute and do not automatically enlarge this legacy shim. +LEGACY_TRAJECTORY_EXPORTS = ( + "differentiate_positions", + "interpolate_with_distance", + "interpolate_with_nums", + "resample_in_time", + "resample_with_distance", + "retime_to_control_grid", + "sort_and_padding_key_frame", + "warp_trajectory_qpos", +) -@pytest.mark.parametrize("name", trajectory.__all__) + +@pytest.mark.parametrize("name", LEGACY_TRAJECTORY_EXPORTS) def test_legacy_trajectory_exports_are_compute_functions(name: str) -> None: assert getattr(action_utils, name) is getattr(trajectory, name)