Add trajectory variants between fixed waypoints

[Yuan-Xinyi]

Description

Affordance sampling (#644) varies where the robot makes contact. This PR covers the other half: once a plan'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. A sampled affordance 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.

It follows #644's structure: the algorithms live at the simulation layer, a tutorial is the demonstration host, and Task Program, Gym lifecycle, dataset persistence and GenerationSession collection bookkeeping are explicitly deferred. No embodichain_tasks/ file is touched.

Refs #644

Main changes

Operators (expansion/operators.py)

• via_points routes an allowed free phase through 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. The knots' signed magnitudes vary along one sampled joint direction: sampling every joint of every knot independently decorrelates the joints, and the tool path that forward kinematics produces then wanders — measured at 1.51x the reference arc length on the Place reference, against 1.11x once the knots share a direction.
• 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. 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.
• Bug fix surfaced by running this: joint-limit rejection previously scanned the whole trajectory, including joints no operator touches. The tutorial's reference holds gripper_finger2_joint_1 at −8e−6, just under its 0.0 lower bound, which rejected 23 of 24 proposals while reporting "Sampled residual violates joint limits". The check now covers only the joints an operator actually moved. This also affected the pre-existing joint_residual.

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)

• spatial.method names one or more joint-path methods. It was a single-choice string, so "enumerate every joint-path operator" could not be expressed at all. It now accepts a sequence, each requested method becomes its own variant, and a bare string is still accepted where a sequence is expected, so the existing spelling keeps working.
• One entry point, one number. expand_trajectory_variants(template, case=..., count=8, joint_limits=..., control_dt=...) is the whole interface, with every implemented method enabled by default. Everything else is an advanced option: default_variant_factors to read or adjust the policy, cfg= for exact control, plan_trajectory_variants / apply_trajectory_variant to separate enumeration from application, and the optional motion-limit and attempt-budget arguments. The tutorial mirrors the split: --trajectory_variants is the only variant flag an ordinary run needs, and the remaining eight sit in an advanced variant options group so --help shows the simple path first. With no advanced flag set the tutorial calls the helpers with no configuration at all, so the documented simple path is the one actually exercised.
• Configuration is optional. expand_trajectory_variants(template, case=..., count=6, ...) is the whole API for "give me six different ways to run this". Omitting cfg resolves default_variant_factors, which enables every implemented factor the supplied inputs support at the magnitudes measured below, and drops the null-space factor when no Jacobians are given rather than proposing variants that would be rejected. An explicit cfg is never overridden, and an explicitly enabled factor that cannot produce a variant now raises instead of disappearing. That covers ik without task_jacobians, and approach, whose Cartesian standoff poses have to be replanned through IK first: the configuration schema still accepts the factor so an upstream planning stage can declare it, but the qpos entry points refuse it rather than silently producing no approach variation. The resolved settings come back on TrajectoryVariantSet.cfg, so an auto-configured run stays reproducible.
• 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.
• apply_trajectory_variant applies one combination, running at most one joint-path operator so a null-space projection is never stacked on an already displaced path.
• 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 variant. Each rejection is counted under a key naming its reason, so a caller can retune the offending setting instead of guessing.
• 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.

Tutorial hook (scripts/tutorials/atomic_action/place.py, tutorial_utils.py)

No new tutorial. Following #644, the shared logic lives in tutorial_utils.py and the existing Place tutorial grows a small hook (+40/-3): --trajectory_variants maps one variant per simulation row, the same shape --affordance_branches uses. Row zero always replays the plan unchanged.

The tutorial declares only which of the action's own named TrajectorySegment ranges may move — gripper-close and release are contact phases, and only motion at or after the lift is retimed so the shared clear_dynamics() step index stays aligned. Everything else comes from the default policy. Advanced overrides sit in their own advanced variant options argument group.

--variant_plot_dir <dir> writes two figures:

• Joint trajectories. Every variant's arm joints against time. Curves separate between waypoints and rejoin inside the shaded contact phases. The title carries the largest contact-phase deviation of any variant from the reference — 0.00e+00 rad, so contact windows are preserved exactly rather than approximately.
• Tool paths. Every variant's tool-centre path in the arena, beside the arm drawn at its starting configuration and the cube drawn to scale on the ground. Stars mark the annotated waypoints, and every variant draws its own, so a waypoint an operator had moved would show up as a scattered cluster instead of a single star; the title reports how far apart they land, which is 0.0e+00 m, together with the lowest tool centre (2.8 cm) so ground clearance is stated rather than implied by a viewing angle.

A junction handler removes the zero-duration sample each concatenated action repeats at a join, after verifying it is a duplicate; a join that actually moves the robot is refused rather than discarded.

What this does and does not guarantee

• 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 a top-down grasp, which is why the tutorial keeps rows (0, 1, 2, 3, 4). Declaring controlled_joint_indices in solver order keeps the columns aligned without any permutation.
• 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.
• The tutorial's lift check is an observation, not a task-success contract.

Figures

Both are committed under docs/source/_static/trajectory_variants/ and embedded in the overview page, so they survive the branch. Produced by the command below with a hundred variants.

Arm joints of 100 variants against time. Curves separate between waypoints and rejoin inside the shaded contact phases. The title carries the largest contact-phase deviation from the reference: 0.00e+00 rad. The wide green band on joint6 is nullspace_residual — for a grasp taken from directly above, wrist roll is the freedom the declared task rows leave open.

Tool paths of the same 100 variants. Stars are the annotated waypoints; every variant draws its own, so they would scatter if an operator had moved one. They spread by 0.0e+00 m. The arm is drawn at its starting configuration and the cube to scale on the ground.

Measured over those hundred: all accepted, no rejected proposals, a hundred distinct geometry families, median tool-path length 1.09x the reference and 1.38x at worst, and three variants whose lowest tool centre dips under 2 cm against a median of 2.8 cm. Sampling more variants reaches further into those tails, which argues for a host-side clearance or collision check rather than trusting the generator alone.

Calling it

Configuration is optional; every implemented factor is enabled by default.

from embodichain.lab.sim.motion.expansion import expand_trajectory_variants

result = expand_trajectory_variants(
    template,                                  # reference with annotated phases
    case=case,                                 # scene and initial-state identity
    count=100,                                 # how many variants to collect
    joint_limits=robot.get_qpos_limits()[0],
    control_dt=0.01,
    task_jacobians=jacobians,                  # optional; unlocks posture variants
)
for variant, row in zip(result.variants, range(len(result.variants))):
    print(variant.spatial_operator, variant.duration_scale, variant.timing_profile)
    positions = result.candidates.positions[row][: result.candidates.valid_length[row]]

The tutorial does exactly this behind one flag:

python scripts/tutorials/atomic_action/place.py --headless --num_envs 1 --device cuda \
  --trajectory_variants 100 --variant_plot_dir outputs/variants

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

Automated checks:

Check	Result
black --check ./	1119 files pass
python docs/scripts/check_api_docs.py	2197/2197 exports documented
context.py check and routing	ok
pytest tests/sim/motion/expansion	195 passed
py_compile on the touched and neighbouring tutorials	passed
pytest tests/test_agent_context_*.py tests/docs	299 passed, 1 failed
python -m py_compile on the tutorial	passed

The one tests/docs failure is test_project_myst_configuration_remains_valid, which raises ModuleNotFoundError: No module named 'myst_parser'. Sphinx and MyST are not installed in this development environment, so that test and the Sphinx docs build could not run here; both are unrelated to this change and covered by the CI build job.

The tutorial was launched on a GPU (RTX 4090) and its results are below, unlike the affordance sampling tutorials in #644, which that PR could not launch:

python scripts/tutorials/atomic_action/place.py --num_envs 6 --device cuda --variant_plot_dir outputs/variant_plots

Trajectory variant 0: operator=none,               scale=1.0,  uniform, samples=239
Trajectory variant 1: operator=joint_residual,     scale=1.0,  uniform, samples=239
Trajectory variant 2: operator=via_points,         scale=1.0,  uniform, samples=239
Trajectory variant 3: operator=nullspace_residual, scale=1.0,  uniform, samples=239
Trajectory variant 4: operator=joint_residual,     scale=1.25, uniform, samples=277
Trajectory variant 5: operator=via_points,         scale=1.25, uniform, samples=277

max contact-phase deviation from the reference: 0.00e+00 rad

Six distinct geometry families covering all three joint-path operators, no rejected proposals, and exact preservation of every contact phase. The two figures are written to --variant_plot_dir; they are not committed, since outputs/ is gitignored. This run exercised the real UR5 solver Jacobian through nullspace_residual, which unit tests alone could not.

Measured tuning

joint_offset_scale is a fraction of each joint's declared range, and this arm declares ±2π per joint, so small fractions are large in absolute terms:

joint_offset_scale	Rows lifting the cube
0.005	6 of 6
0.01 (tutorial default)	6 of 6
0.02	5 of 6
0.03	5 of 6

The usable window is narrow on a task this constrained, and it interacts with deduplication: at the shared default of 0.01 for both joint_offset_scale and coverage.joint_dedup_normalized_tol, the via_points variant collapses into the reference's geometry family. The tutorial therefore defaults --dedup_tolerance to 0.004. Both numbers are documented rather than silently tuned.

Default configuration

The schema's own defaults leave every factor off, set joint_offset_scale to 0.05, via_count to 1 and target_per_cell to 1. Enabling factors one at a time from those values reproduces every problem measured above: a silent single-variant result, offsets past the point where the grasp survives, via_points degenerating to a single bump, and timing variants deduplicated away. Those schema defaults are left untouched, since an all-off configuration is a meaningful explicit state; the working values live in default_variant_factors and apply when a caller does not supply a configuration at all.

Naming

An earlier revision 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, module, docs and tests use 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 production-contract tests that prove the feature works
• Dependencies have been updated, if applicable (no dependency changes required)

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 2/5

The PR does not appear safe to merge yet because three outstanding behavioral defects can break documented APIs or leak generation capacity, and an explicit public-export requirement remains unmet.

Findings

1.  Default variants cannot be applied ▶
2.  Approach factor silently ignored ▶
3.  Invalid Evidence Leaks Capacity ▶
4.  Timing Variants Lose Labels ▶
5.  Timing Scalars Bypass Validation ▶
6.  Public export is missing ▶

Fix with agent prompt

### Issue 1
embodichain/lab/sim/motion/expansion/variants.py:undefined-247
Calling `plan_trajectory_variants(count)` without a configuration enables `nullspace_residual`, but calling `apply_trajectory_variant(...)` with its defaults provides no task Jacobians. Applying one of those planned null-space variants therefore raises `ValueError`, so the documented default plan-then-apply workflow fails. Planning must omit redundancy when its required input is unavailable, or application must receive the resolved configuration and inputs.

### Issue 2
embodichain/lab/sim/motion/expansion/variants.py:249-258
A caller can pass an approach-enabled configuration directly to the public `apply_trajectory_variant` entry point. Unlike planning and expansion, this path neither applies nor rejects the approach factor, so it returns a qpos trajectory without the requested approach variation. Apply the same fail-fast validation in the direct application path.

### Issue 3
embodichain/lab/sim/motion/expansion/session.py:533-544
When banded coverage is enabled, a missing, malformed, negative, or non-finite `manipulability` observation raises during coverage classification without releasing the running attempt. Because running attempts retain their episode-byte reservation, one bad rollout can permanently consume pending capacity and block later candidates. Release the attempt before propagating this evidence-validation failure, as is already done for other rejected episodes.

### Issue 4
scripts/tutorials/atomic_action/tutorial_utils.py:846-856
When more than twelve variants are plotted, this code groups them only by `spatial_operator`. In a timing-only expansion, the reference claims the `none` group first, so variants with different duration scales or timing profiles receive no legend label. The resulting validation figures do not explain what the remaining black curves represent.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

### Issue 5
embodichain/lab/sim/motion/expansion/cfg.py:104-105
This scalar-to-tuple conversion applies to every `tuple[str, ...]` field, not only `spatial.method`. As a result, a malformed value such as `timing.profiles: "uniform"` is converted into a tuple and accepted even though timing profiles require a nonempty sequence. Restricting this compatibility behavior to `spatial.method` would preserve strict validation for timing profiles.

### Issue 6
embodichain/lab/sim/motion/expansion/operators.py:36-44
`TIMING_PROFILES` is documented and re-exported as public API, but the module that defines it omits it from `__all__`. This violates the repository directive that public modules define their public exports, so the requirement must be satisfied before merging.

```suggestion
__all__ = [
    "TIMING_PROFILES",
    "rotate_grasp_about_object_axis",
    "perturb_approach_direction",
    "joint_residual",
    "via_points",
    "nullspace_residual",
    "retime",
    "validate_motion_limits",
]
```

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Summary

The PR adds trajectory-variant generation for fixed waypoints, including spatial, null-space, timing, and approach operators, plus tutorial integration, plotting, documentation, and focused tests.

• Adds deterministic planning and application APIs with coverage-based deduplication.
• Expands configuration and public APIs for spatial, redundancy, approach, and timing factors.
• Integrates variant generation and visualization into the Place tutorial.
• The latest changes decouple generated variant count from replay environment count and improve crowded plot rendering.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[Reference trajectory] --> B[Plan factor combinations]
    B --> C[Apply one path operator]
    C --> D[Apply optional retiming]
    D --> E[Motion-limit checks]
    E --> F[Geometry and timing deduplication]
    F --> G[Accepted variant set]
    G --> H[Plot all accepted variants]
    G --> I[Select or cycle rows for simulation environments]

Loading

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

[greptile-apps]

Default variants cannot be applied

Calling plan_trajectory_variants(count) without a configuration enables nullspace_residual, but calling apply_trajectory_variant(...) with its defaults provides no task Jacobians. Applying one of those planned null-space variants therefore raises ValueError, so the documented default plan-then-apply workflow fails. Planning must omit redundancy when its required input is unavailable, or application must receive the resolved configuration and inputs.

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/lab/sim/motion/expansion/variants.py
Line: 246

Comment:
**Default variants cannot be applied**

Calling `plan_trajectory_variants(count)` without a configuration enables `nullspace_residual`, but calling `apply_trajectory_variant(...)` with its defaults provides no task Jacobians. Applying one of those planned null-space variants therefore raises `ValueError`, so the documented default plan-then-apply workflow fails. Planning must omit redundancy when its required input is unavailable, or application must receive the resolved configuration and inputs.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[greptile-apps]

Timing Scalars Bypass Validation

This scalar-to-tuple conversion applies to every tuple[str, ...] field, not only spatial.method. As a result, a malformed value such as timing.profiles: "uniform" is converted into a tuple and accepted even though timing profiles require a nonempty sequence. Restricting this compatibility behavior to spatial.method would preserve strict validation for timing profiles.

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/lab/sim/motion/expansion/cfg.py
Line: 104-105

Comment:
**Timing Scalars Bypass Validation**

This scalar-to-tuple conversion applies to every `tuple[str, ...]` field, not only `spatial.method`. As a result, a malformed value such as `timing.profiles: "uniform"` is converted into a tuple and accepted even though timing profiles require a nonempty sequence. Restricting this compatibility behavior to `spatial.method` would preserve strict validation for timing profiles.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[greptile-apps]

Approach factor silently ignored

A caller can pass an approach-enabled configuration directly to the public apply_trajectory_variant entry point. Unlike planning and expansion, this path neither applies nor rejects the approach factor, so it returns a qpos trajectory without the requested approach variation. Apply the same fail-fast validation in the direct application path.

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/lab/sim/motion/expansion/variants.py
Line: 249-258

Comment:
**Approach factor silently ignored**

A caller can pass an approach-enabled configuration directly to the public `apply_trajectory_variant` entry point. Unlike planning and expansion, this path neither applies nor rejects the approach factor, so it returns a qpos trajectory without the requested approach variation. Apply the same fail-fast validation in the direct application path.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[greptile-apps]

Invalid Evidence Leaks Capacity

When banded coverage is enabled, a missing, malformed, negative, or non-finite manipulability observation raises during coverage classification without releasing the running attempt. Because running attempts retain their episode-byte reservation, one bad rollout can permanently consume pending capacity and block later candidates. Release the attempt before propagating this evidence-validation failure, as is already done for other rejected episodes.

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/lab/sim/motion/expansion/session.py
Line: 533-544

Comment:
**Invalid Evidence Leaks Capacity**

When banded coverage is enabled, a missing, malformed, negative, or non-finite `manipulability` observation raises during coverage classification without releasing the running attempt. Because running attempts retain their episode-byte reservation, one bad rollout can permanently consume pending capacity and block later candidates. Release the attempt before propagating this evidence-validation failure, as is already done for other rejected episodes.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[greptile-apps]

Public export is missing

TIMING_PROFILES is documented and re-exported as public API, but the module that defines it omits it from __all__. This violates the repository directive that public modules define their public exports, so the requirement must be satisfied before merging.

Suggested change

	__all__ = [
	"rotate_grasp_about_object_axis",
	"perturb_approach_direction",
	"joint_residual",
	"via_points",
	"nullspace_residual",
	"retime",
	"validate_motion_limits",
	]
	__all__ = [
	"TIMING_PROFILES",
	"rotate_grasp_about_object_axis",
	"perturb_approach_direction",
	"joint_residual",
	"via_points",
	"nullspace_residual",
	"retime",
	"validate_motion_limits",
	]

Context Used: CLAUDE.md (source)

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/lab/sim/motion/expansion/operators.py
Line: 36-44

Comment:
**Public export is missing**

`TIMING_PROFILES` is documented and re-exported as public API, but the module that defines it omits it from `__all__`. This violates the repository directive that public modules define their public exports, so the requirement must be satisfied before merging.

```suggestion
__all__ = [
    "TIMING_PROFILES",
    "rotate_grasp_about_object_axis",
    "perturb_approach_direction",
    "joint_residual",
    "via_points",
    "nullspace_residual",
    "retime",
    "validate_motion_limits",
]
```

**Context Used:** CLAUDE.md ([source](https://github.com/dexforce/embodichain/blob/main/CLAUDE.md))

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[greptile-apps]

Timing Variants Lose Labels

When more than twelve variants are plotted, this code groups them only by spatial_operator. In a timing-only expansion, the reference claims the none group first, so variants with different duration scales or timing profiles receive no legend label. The resulting validation figures do not explain what the remaining black curves represent.

Prompt To Fix With AI

This is a comment left during a code review.
Path: scripts/tutorials/atomic_action/tutorial_utils.py
Line: 846-856

Comment:
**Timing Variants Lose Labels**

When more than twelve variants are plotted, this code groups them only by `spatial_operator`. In a timing-only expansion, the reference claims the `none` group first, so variants with different duration scales or timing profiles receive no legend label. The resulting validation figures do not explain what the remaining black curves represent.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[yuecideng]

[P1] Apply ik.task_rows or remove the public knob — This field is validated and documented as a configuration option, but no code in the variant application path reads it; the full task_jacobians tensor is passed straight to nullspace_residual. A caller selecting a subset of spatial rows therefore gets the same projection as the default, with no warning. Select these rows at one clearly owned boundary (and test a full six-row input with a custom subset), or remove the field and document that callers must pre-reduce the Jacobian.

[yuecideng]

[P1] Do not swallow malformed inputs as proposal rejections — apply_trajectory_variant raises ValueError for structural or caller errors as well as for an operator rejecting a sampled proposal (for example invalid Jacobian or joint-limit shapes, or an invalid control period). Catching every ValueError here converts those programming/configuration errors into rejected-count bookkeeping and can return a partial result after the nominal variant has already succeeded. Prevalidate inputs before the attempt loop or introduce a dedicated rejection exception and catch only that type.

[yuecideng]

[P2] Use a numerical rank test instead of largest <= 0 — For a generic full-rank, non-orthogonal Jacobian, I - pinv(J) @ J contains floating-point residuals even though the null space is mathematically zero, so largest will usually be positive. The current check therefore only reliably rejects special cases such as an identity Jacobian and may emit a near-unchanged variant when no redundancy exists. Use the singular values/rank with the existing rank_tolerance, or compare the projector effect against a scale-aware tolerance, and add a non-diagonal full-rank regression test.