Add geometry-constrained parallel affordance expansion

[matafela]

Description

Add --n_affordance_expand N to generate and execute affordance variations across N parallel environments. Variations follow task geometry, using valid grasp candidates, declared symmetries, or accepted motion ranges.

Main changes

• Introduce AffordanceSamplingContext and AffordancePoseCandidates for reproducible sampling, explicit candidate validity, and diversity-based selection. Preserve the nominal branch and report candidate reuse when distinct alternatives are exhausted.
• Connect candidate selection to action planning: (B, K, 4, 4) candidate poses become one selected pose per environment, then action-specific waypoints and (B, T, D) joint trajectories. affordance_sampling.py handles selection; action planners retain ownership of IK and trajectory generation.
• Preserve PickUp’s feasibility checks and each selected grasp’s object-to-EEF transform for subsequent transport and placement.
• Support constrained variations for grasping, axis alignment, pressing, door opening, sliding, twisting, interaction points, and assembly symmetries. Missing optional motion ranges preserve exact task targets.
• Count max_episodes as accepted per-environment episodes and consistently commit partial final batches across dataset, trajectory, and camera recording.
• Add a temporary Default/CUDA drawer-handle material refresh at startup to address first-environment grasp slip. The refresh restores the original friction immediately and leaves mass, inertia, and joint state unchanged.
• This change also updates the shared UR5 gripper grasp target from 0.024 to 0.040, the shared trajectory sample count from 40 to 100, and drawer hand interpolation from 12 to 18. Shared component changes affect all deployments referencing those components.

Example commands

conda activate embodichain2

embodichain run-env \
  --gym_config embodichain_tasks/configs/tasks/manipulation/repeated_pick_place/task.ur5.yaml \
  --device cuda --headless --seed 42 \
  --max_episodes 180 --n_affordance_expand 9

embodichain run-env \
  --gym_config embodichain_tasks/configs/tasks/manipulation/open_drawer/task.ur5.yaml \
  --device cuda --headless --seed 42 \
  --max_episodes 180 --n_affordance_expand 9

Nine branches mean nine total environments. With all rows accepted, 20 episodes require three batches (9 + 9 + 2), while 180 episodes require twenty batches.

Append --record_trajectory --trajectory_save_dir ./outputs/affordance_trajectories to save joint trajectories.

Documentation

• docs/source/overview/sim/atomic_actions/affordance_expansion.md: commands, sampling algorithm, candidate-to-trajectory pipeline, geometric constraints, collection semantics, and drawer tuning/workaround.
• docs/source/overview/sim/atomic_actions/builtin_actions.md: updated action behavior and optional sampling ranges.
• docs/source/api_reference/public_api.rst: public sampling API.
• agent_context/topics/atomic-actions/affordance-expansion.md: implementation ownership and integration contracts.

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)

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.

[greptile-apps]

Confidence Score: 4/5

The PR is not safe to merge with camera recording enabled because mixed or partial expanded batches can produce incorrect videos, omit committed environments, or fail recorder finalization.

Findings

1.  Camera Subset Commits Break ▶

Fix with agent prompt

### Issue 1
embodichain/lab/gym/envs/embodied_env.py:993-1001
When an expanded batch commits only some environments, this code passes that subset to every camera recorder. The synchronous recorder ignores `env_ids` and saves the full multi-environment frame mosaic, so failed or unselected rows are included. The asynchronous recorder tracks at most four environments and requires every tracked row to have a committed episode before flushing, so selected rows numbered four or higher can be dropped and partial commits can fail during finalization. Camera recording therefore does not match the dataset and trajectory commit boundary.

---

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

Summary

This PR adds reproducible, geometry-constrained affordance expansion across parallel simulation environments and wires it through Atomic Skills, Task Programs, collection, and several new task deployments.

• Adds candidate packing, validity masking, deterministic branch selection, range sampling, and diagnostic metadata.
• Integrates sampled grasps and motion ranges into PickUp, Press, Slide, Twist, OpenDoor, AxisAlign, and assembly planning.
• Adds configured Press, Twist, OpenDoor, and Pour deployments plus drawer contact and trajectory tuning.
• Updates accepted-row collection so dataset, trajectory, and camera outputs are intended to share a partial commit boundary; camera recorders do not yet satisfy that boundary.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Observed scene and task geometry] --> B[Generate valid pose candidates]
    B --> C[AffordancePoseCandidates B × K]
    C --> D[Deterministic branch selection]
    D --> E[One selected pose per environment]
    E --> F[Action-specific IK and trajectory planning]
    F --> G[Parallel environment execution]
    G --> H{Successful row with frames?}
    H -- Yes --> I[commit_env_ids subset]
    H -- No --> J[Discard row during reset]
    I --> K[Dataset save]
    I --> L[Trajectory save]
    I --> M[Camera save]
    M --> N[Current recorder APIs cannot preserve arbitrary subsets]

Loading

_{Reviews (1) · Last reviewed commit: "update"}

[greptile-apps]

Camera subset commits break

When an expanded batch commits only some environments, this code passes that subset to every camera recorder. The synchronous recorder ignores env_ids and saves the full multi-environment frame mosaic, so failed or unselected rows are included. The asynchronous recorder tracks at most four environments and requires every tracked row to have a committed episode before flushing, so selected rows numbered four or higher can be dropped and partial commits can fail during finalization. Camera recording therefore does not match the dataset and trajectory commit boundary.

Prompt To Fix With AI

This is a comment left during a code review.
Path: embodichain/lab/gym/envs/embodied_env.py
Line: 993-1001

Comment:
**Camera subset commits break**

When an expanded batch commits only some environments, this code passes that subset to every camera recorder. The synchronous recorder ignores `env_ids` and saves the full multi-environment frame mosaic, so failed or unselected rows are included. The asynchronous recorder tracks at most four environments and requires every tracked row to have a committed episode before flushing, so selected rows numbered four or higher can be dropped and partial commits can fail during finalization. Camera recording therefore does not match the dataset and trajectory commit boundary.

---

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

[yuecideng]

The refresh_articulation_contact_material workaround should not be implemented at the EmbodiChain event layer. The stale contact-material binding is a DexSim backend initialization defect and should be fixed at the owning native material-binding boundary.

Temporarily perturbing dynamic_friction and restoring it through private native handles couples a packaged task to DexSim internals, bypasses the public articulation abstraction, and can hide the underlying initialization failure. It also makes correctness depend on a task-specific startup event being present.

Please fix the material binding in DexSim so the authored contact material is valid before the first physics step, then remove _refresh_link_contact_material, refresh_articulation_contact_material, the open_drawer startup event, and their workaround-specific tests/documentation from this PR.

[yuecideng]

Documentation-focused review: I found two medium-priority contract/reproducibility issues and one low-priority observability mismatch. Details are attached inline.

[yuecideng]

[P2] This overstates the camera recorder contract. The synchronous record_camera_data.save_and_clear() ignores env_ids and persists a grid that already contains every environment, so rejected rows remain in the video. The asynchronous recorder queues only the selected rows, but _flush_committed_episodes() requires every recorder-local queue to be non-empty and finalize() rejects an incomplete partial batch. Therefore, with mixed-success rows or a partial final batch, camera output does not honor the same subset as dataset/trajectory output and may either contain uncommitted rows or fail at finalization. Please either implement true row-selective camera commits or document the two recorder limitations instead of claiming subset parity. See embodichain/lab/gym/envs/managers/record.py lines 140-158, 193-200, and 226-302.

[yuecideng]

[P2] The runnable example depends on an undocumented local environment name. The repository never instructs users to create a Conda environment named embodichain2; the public installation guide uses a .venv workflow. A clean user following the installation docs will fail immediately at conda activate embodichain2. Please make this environment-agnostic (for example, 'run in the environment where EmbodiChain is installed') and link the installation guide, or document the complete Conda setup.

[yuecideng]

[P3] This reads as if the listed candidate diagnostics are retained by every expanded action plan, but several documented paths discard them. AxisAlign._resolve_grasp_pose() binds the third value from select() to _ and its build_plan() call does not add selection metadata; InteractionPoints.sample_poses() and AssembleAffordance.get_assemble_object_pose() similarly return only select(...)[1]. Those paths retain only the generic expansion-context metadata, not candidate_ids, valid/unique counts, or reused. Please scope this statement to the actions that actually retain selection metadata, or propagate the metadata through these paths.