Add trajectory variants between fixed waypoints

[Yuan-Xinyi]

Description

Affordance expansion (#637) varies where the robot makes contact. This PR covers the other half: once a task's waypoints are settled, produce several different ways to execute them, so downstream imitation learning and reinforcement-learning post-training see more than one solution per task instead of many copies of one.

The two compose. An affordance variation produces a new set of waypoints, and variant expansion then produces several ways of executing that set.

This also gives motion/expansion/ its first caller. The package already shipped candidate contracts, coverage bookkeeping and GenerationSession, but nothing in the repository used them, and only two of its seven declared augmentation factors were implemented.

Main changes

Operators (expansion/operators.py)

• via_points routes an allowed free phase through sampled interior knots, interpolated with clamped cubic Hermite segments. joint_residual adds one fixed-shape bump per phase; with two or more knots this changes the shape of the path, not only its amplitude.
• nullspace_residual projects a residual onto the null space of task Jacobians supplied by the caller, changing arm posture while holding the declared task rows. A fully constrained task raises rather than silently returning the reference.
• retime gains a bounded within-phase profile (uniform, ease_in, ease_out) that redistributes time inside a phase without changing the path or the phase's total duration, so one path yields several velocity profiles. The uniform path is arithmetically unchanged from the existing operator.
• perturb_approach_direction places standoff poses on a cone around a nominal approach direction while leaving the contact transform exact.

Every qpos operator uses an envelope that is zero, with zero derivative, at both endpoints of the phase it modifies, and contact and hold phases are never touched. Annotated waypoints, contact windows and dwell durations stay bit-identical to the reference.

Variant generation (expansion/variants.py, new)

• plan_trajectory_variants lists the enabled factor combinations. Ordinal zero is always the unmodified reference; later ordinals cycle through the enabled joint-path operators, then the duration scales, then the time warps.
• expand_trajectory_variants collects variants for one fixed scene, rejecting proposals that an operator refuses, that fail sampled motion limits, or whose measured geometry and timing duplicate an accepted row.
• expand_row_variants assigns one variant per independent execution row, meaning one per parallel environment with its own randomized scene. Nothing is deduplicated across rows because geometries from different scenes are not comparable, and a rejected row falls back to its own reference rather than leaving an environment without a command stream. Its ordinal_offset lets repeated rollouts start further along the list instead of repeating the first combinations.
• Every rejection is counted and reported rather than dropped silently.
• CoverageIndex.family_of exposes the resolved geometry family, so geometry_family_id reflects measured grouping rather than a fresh digest.

Configuration (expansion/cfg.py)

ik and approach were _DisabledFactorCfg stubs that raised when enabled; they are now real settings. spatial gains via_count and timing gains profiles. contact, contact_timing and recovery remain unimplemented and are still rejected when enabled.

Packaged host integration

StackBlocksTwo-v1 is opt-in through extensions.trajectory_variants and ships disabled, so its default behavior is unchanged. Two points worth review attention:

• Phase boundaries come from the atomic actions' own named TrajectorySegment ranges rather than from re-deriving their internal step arithmetic. A guard fails loudly if the compiled close segment stops lining up with _insert_grasp_hold's insertion index, instead of silently mis-annotating a contact phase.
• Retiming is limited to phases at or after the lift. The task's executor clears the held block's dynamics at one shared step index, so changing how long the approach takes would put that event at a different step in each environment. Path and posture variation still apply to the approach, because they do not change the sample count.

Each demonstration segment records the episode index, the variant assigned to every row, per-row sample lengths, per-row resolved phase boundaries, and fallback counts under metadata["trajectory_variants"].

What this does and does not guarantee

Stated explicitly because these are easy to overstate:

• nullspace_residual holds the declared task rows to first order only. Phase endpoints stay exact because the envelope vanishes there, but interior samples drift with linearization error and need forward-kinematics verification by the host.
• The Jacobian's reference frame and column order are the caller's responsibility. BaseSolver.get_jacobian returns a base-frame Jacobian, so dropping its angular-z row removes rotation about base z, not about the tool axis. Those coincide for the stack-blocks top-down grasp, which is why it declares task_rows: [0, 1, 2, 3, 4]; a side-grasping task must choose differently.
• Deduplication compares measured joint geometry and elapsed phase time. It is a similarity measure, not a validity proof. Nothing here establishes collision freedom, dynamic feasibility or task success.
• GenerationSession's episode budgets, commit receipts and durable persistence are not used by this path. Persistence stays with the environment's own dataset manager.

Example commands

Run these in the environment where EmbodiChain is installed. Set enabled to true in embodichain_tasks/configs/tasks/manipulation/tableware/stack_blocks_two/env.json, then:

embodichain run-env \
  --gym_config embodichain_tasks/configs/tasks/manipulation/tableware/stack_blocks_two/env.json \
  --num_envs 8 --device cuda --headless --seed 42 --max_episodes 24

Each parallel environment receives a different variant, so --num_envs sets how many variants one rollout collects. Append --record_trajectory --trajectory_save_dir ./outputs/trajectory_variants to save the joint trajectories.

Documentation

• docs/source/overview/sim/motion/trajectory_variants.md: the factor table, what is and is not guaranteed, the runnable example, how to choose between the APIs, the Jacobian frame caveat, and the limitations.
• docs/source/api_reference/embodichain/embodichain.lab.sim.motion.expansion.rst: the new operators, the variants section, and corrected configuration prose.
• agent_context/topics/motion-planning/motion-planning.md and agent_context/MAP.yaml: ownership, boundaries and routing keywords.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• Enhancement (non-breaking change which improves an existing functionality)
• New feature (non-breaking change which adds functionality)
• Breaking change (existing functionality will not work without user modification)
• Documentation update

Validation

Check	Result
black --check ./	1052 files pass
python docs/scripts/check_api_docs.py	2098/2098 exports documented
context.py check and routing	ok
pytest tests/sim/motion (full, 36 min)	740 passed, 37 skipped
pytest tests/sim/motion/expansion	187 passed
pytest tests/gym/envs/test_stack_blocks_two_trajectory_variants.py	9 passed
pytest tests/gym/envs/test_official_task_layout.py tests/gym/utils/test_gym_utils.py	128 passed
pytest tests/test_agent_context_*.py tests/docs	47 passed

Rebased onto main after #648 and #650 landed. #648 rewrote agent_context/topics/motion-planning/motion-planning.md into a terser "find the owner" format specifically to prevent PR-driven content growth, so the context update here was re-resolved onto that format as a minimal delta (+17 lines) rather than the longer section originally written. Neither #648 nor #650 touches Python runtime code, so the tests/sim/motion result above still applies; the gates and affected tests were re-run after the rebase and again after the rename described below.

Not run: the Sphinx docs build. Sphinx is not installed in the development environment used here, so the new page's cross-references were checked for syntax and relative paths only. The CI build job covers it.

The task integration is tested without a simulator: the tests drive the task's own planning helpers with stand-in compiled segments, pinning where fixed phases land, that contacts and dwells survive augmentation bit-identically, that shorter rows hold their final command, that successive episodes move further along the combination list, and that the solver Jacobian columns are permuted into the task's joint order.

Naming

An earlier revision of this branch called these "trajectory modes". That word was invented for this change; the expansion package already says "variant" (CoverageIndex documents "timing variants per geometry") and #637 says "variations". The API, configuration key, module, docs and tests were renamed to variant so the change reuses the repository's existing vocabulary instead of adding a parallel term.

Checklist

• I have run the black . command to format the code base.
• I have made corresponding changes to the documentation
• Public API changes are reflected in the API docs (python docs/scripts/check_api_docs.py), if applicable
• I have added tests that prove my fix is effective or that my feature works
• Dependencies have been updated, if applicable.

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge because the changes since the previous review consistently complete the terminology rename without introducing a new actionable regression.

Summary

This PR adds configurable trajectory variants for fixed waypoints and integrates them as an opt-in feature for StackBlocksTwo-v1.

• Adds joint-path, null-space posture, approach-direction, duration, and timing-profile variation.
• Adds variant planning, application, deduplication, fallback accounting, and per-row metadata.
• Integrates phase-aware expansion into the stack-blocks task while preserving contact and hold phases.
• Renames the newly introduced “trajectory modes” API and configuration vocabulary to “trajectory variants.”
• Adds focused operator, configuration, expansion, task-integration, and documentation coverage.

_{Reviews (2) · Last reviewed commit: "feat(motion): add trajectory variants be..."}