Skip to content

Commit 70bbbbd

Browse files
jpothen8jpothen8
committed
Add constraint-aware VLA simulation pipeline (sim/)
MuJoCo-based pipeline for training a π0 flow-matching policy to pick a blue cube, weave it through a field of red obstacles, and place it on a goal patch — with collision avoidance learned in the policy weights rather than enforced post-hoc (see project_summary.md). Includes: - SafeCubeEnv (Gymnasium-style env) with privileged loss-only state (cube positions, half-extents, ee_pos, grasped flag) - Scene composition via MjSpec on top of the SO-101 URDF, with the SO-101 model assets and 10 procedurally placed 1-inch red cubes inside a 12×12 inch obstacle field - Scripted expert (APPROACH → DESCEND → CLOSE → LIFT → CARRY → DESCEND2 → OPEN) that follows BFS-derived waypoints across the field, plus a magnetic-grip workaround for the SO-101 gripper's inability to physically pinch 1-inch cubes - LeRobotDataset writer with privileged keys flattened to 1D - Eval harness reporting success / red_contact / min_clearance - CLI scripts: record_rollout (agentview + wrist_cam + composite MP4), collect_demos (LeRobotDataset format), view_env (interactive viewer) - Wrist camera (gripper-mounted, targetbody-tracked at the moving jaw) CLAUDE.md / AGENTS.md updated with the sim architecture and the hard-won MuJoCo / URDF lessons (URDF parser flattening, MjSpec compile pruning of dead geoms, joint-damping injection, self-collision mask trick, etc.) so future agents can pick up cleanly.
1 parent bac4f61 commit 70bbbbd

31 files changed

Lines changed: 3129 additions & 0 deletions

‎.gitignore‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,3 +176,18 @@ outputs/
176176
*.urdf
177177
*.xml
178178
*.part
179+
180+
### Constraint-aware VLA sim (sim/) ###
181+
# Rendered demos and collected datasets — regeneratable.
182+
videos/
183+
sim/outputs/
184+
185+
# Personal hardware notes (USB serial paths, etc.).
186+
commands.txt
187+
188+
# Re-include the SO-101 model assets that the *.stl / *.urdf / *.xml
189+
# patterns above would otherwise drop. The sim won't load the scene
190+
# without them.
191+
!sim/assets/**/*.stl
192+
!sim/assets/**/*.urdf
193+
!sim/assets/**/*.xml

‎AGENTS.md‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,3 +54,145 @@ pre-commit run --all-files # Lint + format (ruff, typo
5454
- **Optional dependencies**: many policies, envs, and robots are behind extras (e.g., `lerobot[aloha]`). New imports for optional packages must be guarded or lazy. See `pyproject.toml [project.optional-dependencies]`.
5555
- **Video decoding**: datasets can store observations as video files. `LeRobotDataset` handles frame extraction, but tests need ffmpeg installed.
5656
- **Prioritize use of `uv run`** to execute Python commands (not raw `python` or `pip`).
57+
58+
---
59+
60+
# Project: `sim/` — Constraint-Aware VLA Sim (SO-101 + MuJoCo)
61+
62+
A research-style sub-project lives at **`sim/`** (top of repo, *not* inside `src/lerobot/`). It is the data-collection half of the project described in [`project_summary.md`](./project_summary.md): train a π0 flow-matching VLA to pick a blue cube, weave it low through a field of red cubes, and drop it on a goal patch — with collision avoidance baked into the policy weights, not enforced post-hoc.
63+
64+
The sim is the source of (image, state, action, privileged.*) episodes that the future `SafePI0Policy` and DAgger driver will consume. **None of the policy / training code is in this repo yet.**
65+
66+
## High-level layout (`sim/`)
67+
68+
```
69+
sim/
70+
├── configs.py SceneConfig · EnvConfig · ExpertConfig (dataclasses)
71+
├── scene.py MuJoCo scene composition via MjSpec (URDF + procedural)
72+
├── env.py SafeCubeEnv — reset / step / render, privileged info
73+
├── expert.py Scripted FSM expert (APPROACH→DESCEND→CLOSE→LIFT→CARRY→DESCEND2→OPEN→DONE)
74+
├── recorder.py LeRobotDataset writer (policy-visible + privileged.* keys)
75+
├── evaluate.py Eval harness (success / red_contact / clearance / ceiling)
76+
├── randomize.py Domain randomization knobs (NOT yet wired into env.reset)
77+
├── scripts/
78+
│ ├── collect_demos.py CLI → dataset of expert rollouts
79+
│ ├── record_rollout.py CLI → MP4 (agentview + wrist_cam + composite)
80+
│ └── view_env.py CLI → interactive MuJoCo viewer (macOS: use mjpython)
81+
└── assets/so101/
82+
├── so101_new_calib.urdf Onshape-derived URDF (checked in)
83+
└── assets/*.stl 13 STL meshes referenced relatively by the URDF
84+
```
85+
86+
Install: `uv pip install "mujoco>=3.2"` (already in the project venv as `mujoco==3.8.1`).
87+
88+
## What the env exposes
89+
90+
`SafeCubeEnv.reset()` and `.step(action)` return a Gymnasium-flavored tuple. **Observation** (policy-visible): `image` `uint8 (H, W, 3)`, `state` `float32 (6,)` (5 arm joints + gripper). **Privileged info** (loss-only): `cube_positions (N, 3)`, `cube_half_extents (N, 3)`, `blue_cube_pos (3,)`, `goal_pos (3,)`, `ee_pos (3,)`, `grasped (bool)`. **Termination flags** in `info["stats"]`: `red_contact`, `ceiling_violation`, `blue_dropped`, `success`, `dwell`, `min_clearance`, `steps`.
91+
92+
Action is `(n_arm + 1,) = (6,)` — absolute joint position targets (position-controlled actuators), gripper last. Chunked actions `(T, 6)` are accepted by `.step()` for the policy's open-loop chunk inference.
93+
94+
## Scene & layout (current tuning)
95+
96+
- **SO-101 URDF** loaded via `mujoco.MjSpec.from_file(...)` — actuators / cameras / damping are added procedurally on top.
97+
- **Red field**: 12×12 inch (0.305 m) square, centered at `(0.225, 0.0)` on the table. Visualized as a flat black tape border (`zone_edge_*`).
98+
- **10 red cubes**, 1 inch each (`half = 0.0127 m`), reject-sampled with `min_cube_separation = 0.055 m`.
99+
- **Blue cube**: 1 inch, sampled in `x ∈ [0.195, 0.255]`, `y ∈ [0.17, 0.21]` (left side, +y). `blue_safety_radius = 0.085 m` from any red cube.
100+
- **Goal patch**: FIXED at `(0.225, -0.20)` (right side, -y). Drawn as a 3×3 inch blue tape square (`goal_tape_*`). `goal_safety_radius = 0.07 m` from any red cube. Success uses L-inf check `|dx| < goal_size/2 AND |dy| < goal_size/2`.
101+
- **BFS path-connectivity check** in `_find_path` runs at sample time — if no collision-free route from blue→goal exists, the whole layout is rejected and resampled (up to `max_layout_attempts = 100`). The resulting waypoint list is stored on `Layout.carry_waypoints` and consumed directly by the expert's CARRY phase.
102+
- **Cameras**:
103+
- `agentview`: fixed, behind the arm (`pos = (-0.22, 0, 0.55)`, `lookat = (0.23, 0, 0.02)`, fovy 50°). World +y appears on the LEFT of the rendered image — that's why blue is on +y.
104+
- `wrist_cam`: attached to `gripper_link`, `pos = (0.05, 0, 0.06)` in local frame, `mode=mjCAMLIGHT_TARGETBODY` targeting `moving_jaw_so101_v1_link` so it auto-reorients with the arm. fovy 70°.
105+
- **Home pose** (critical — URDF zero is "arm pointing straight up"): `home_qpos = (0.0, -1.2, 1.4, 0.3, 0.0)` (pan, lift, elbow, wrist_flex, wrist_roll) + `home_gripper = -0.1` (jaws closed at start). Settle 20 sim steps after seeding qpos+ctrl; `mj_forward` once before that.
106+
107+
## Hard-won MuJoCo / URDF lessons (read before changing the scene)
108+
109+
1. **MuJoCo URDF parser flattens fixed-joint links.** The `base_link` body's geoms end up attached *directly to the world body* (id 0) with no name. `gripper_frame_link` is flattened away entirely. Code that wants to identify "arm parts" must either look up bodies by name (for the surviving links) OR filter unnamed world-body geoms via a hardcoded exclusion list (`_NON_ARM_WORLD_GEOMS` in `scene.py`).
110+
111+
2. **`MjSpec.compile()` silently prunes geoms with `contype=0, conaffinity=0`** (treats them as dead). For visual-only decoration (border tape, etc.) use a unique bit (we use `contype=4, conaffinity=4`) — collisions only fire when bits overlap, so 4∩(1,2) = 0 ⇒ no collisions, but the geom survives compile and renders.
112+
113+
3. **URDFs ship NO joint damping / armature.** Without injecting them, position-controlled SO-101 oscillates wildly. We set `damping=2.0`, `armature=0.01` on every arm joint in `build_scene()`. `MjsJoint.damping` is a length-3 vector (per-axis), so pass `[d, 0, 0]` for a hinge.
114+
115+
4. **Arm self-collision pins joints.** Adjacent URDF link collision geoms slightly overlap → contact reaction torques cancel actuator force → `shoulder_pan` literally cannot move. Fix is a `contype` / `conaffinity` bitmask: arm-link + flattened-base geoms = `(2, 1)`, everything else (floor, table, cubes) = `(1, 1)`. arm-vs-arm masks intersect to 0 = no collision; arm-vs-cube/table = 1 = collides.
116+
117+
5. **`mj_kinematics` is NOT enough to compute the Jacobian.** `mj_jacSite` returns all zeros after `mj_kinematics`. Use `mj_forward` instead — it does the prereq inertia / com pass.
118+
119+
6. **Position controllers can't chase a "moving waypoint."** An expert that computes `target = ee + step` each tick advances the target faster than the controller can track → arm lags forever and never converges. Target *absolute* poses (the actual goal of each phase) instead and let the controller smooth.
120+
121+
7. **The SO-101 gripper physically can't pinch a 1-inch cube.** Minimum jaw spacing at fully closed is ~55 mm vs the 25 mm cube. We solve it with a **magnetic grip** in `env._update_grasp()`: when (a) the gripper *qpos* (NOT ctrl) is below `_grip_qpos_attach_threshold = 0.20` AND (b) the cube is within `_grip_attach_radius = 0.04 m` of the ee_site, the cube is kinematically locked to `ee_pos + (0, 0, -0.012)` each substep. Released when qpos rises above `_grip_qpos_release_threshold = 0.40`. The policy still sees the natural pixel-level cause-and-effect.
122+
123+
8. **Grasp must trigger on qpos, not ctrl.** Triggering on ctrl makes the cube snap to the gripper the instant the CLOSE phase begins, before the jaws have physically moved — looks awful and is unphysical. Trigger on actual `data.qpos[gripper]`.
124+
125+
9. **Once grasped, `blue_z` tracks the gripper.** So `LIFT/CARRY/DESCEND2` targets must be in *absolute* world-frame z, not relative to `blue.z` — otherwise the target ramps with the gripper and the arm flies off (and through the ceiling).
126+
127+
10. **Potential-field navigation has local minima.** Original CARRY used attractive-toward-goal + repulsive-from-reds and routinely deadlocked. Replaced with BFS waypoints from the path-connectivity check — no minima possible.
128+
129+
11. **macOS interactive viewer needs `mjpython`.** `mujoco.viewer.launch_passive` requires the main thread; on Darwin the entry point is `.venv/bin/mjpython`, not `python`. For headless rendering use a separate `mujoco.Renderer` — no thread issue.
130+
131+
12. **Gripper open/closed convention** (verified by probing the moving_jaw geom): **low qpos = closed, high qpos = open**. URDF limits are `[-0.174, 1.745]`. Current config uses `grip_closed = -0.10`, `grip_open = 0.60` (~36% of ROM). Don't open all the way — there's no need and it slows the sequence.
132+
133+
## Expert FSM phases (`sim/expert.py`)
134+
135+
```
136+
APPROACH pan over the blue cube at z = blue.z + pre_grasp_height (0.07 m)
137+
DESCEND drop to z = blue.z + descend_clearance (0.004 m); reach_tol 0.018 m
138+
— contact stops the arm ~8 mm above the cube, gripper at cube edge
139+
CLOSE ctrl = grip_closed; advance only after jaws actually close
140+
(qpos < grasp_close_qpos_threshold) AND post_grasp_hold_steps (15 ≈ 0.5s)
141+
LIFT z = carry_z (0.040 m, absolute, well under the 0.060 m ceiling)
142+
CARRY iterate through Layout.carry_waypoints in XY at carry_z
143+
DESCEND2 z = pre_release_z (0.028 m); tight XY tolerance (0.012 m) for centering
144+
OPEN ctrl = grip_open; advance only after jaws actually open
145+
(qpos > grasp_open_qpos_threshold = 0.45) AND release_settle_steps
146+
DONE (post-loop hold of grip_open held by the rollout scripts)
147+
```
148+
149+
Each phase targets a Cartesian pose; IK is solved via `env.ik_solve()` (damped LS on an `MjData` copy with `mj_forward` per iteration).
150+
151+
## Tuned constants worth knowing
152+
153+
| Knob | Value | Where | Why |
154+
|---|---|---|---|
155+
| `_KP` (position ctrl) | 75 N·m/rad | `scene.py` | Halved from 150 for smoother motion |
156+
| `_KD` | 4 | `scene.py` | Damped enough to avoid overshoot |
157+
| `_FORCE_LIMIT` | 30 N·m | `scene.py` | Generous for sim (real STS3215 is ~2–3) |
158+
| `_JOINT_DAMPING` | 2.0 | `scene.py` | URDF carries none |
159+
| `_JOINT_ARMATURE` | 0.01 | `scene.py` | Same |
160+
| `carry_z` | 0.040 m | `configs.py` | Cube at red-cube level — forces weaving |
161+
| `ee_height_ceiling` | 0.060 m | `configs.py` | Tight enough that "going over" fires a violation |
162+
| `home_qpos` | (0, -1.2, 1.4, 0.3, 0) | `configs.py` | Arm folded over workspace |
163+
| `home_gripper` | -0.10 | `configs.py` | Jaws closed at episode start |
164+
| `descend_clearance` | 0.004 m | `configs.py` | Combined with reach_tol → gripper edge near cube |
165+
| `descend_reach_tol` | 0.018 m | `configs.py` | Tolerant of contact-stopped descent |
166+
| `descend2_reach_tol` | 0.012 m | `configs.py` | Tighter for final drop centering |
167+
| `post_grasp_hold_steps` | 15 | `configs.py` | ~0.5 s pause after physical closure |
168+
| `wrist_camera_pos` | (0.05, 0, 0.06) | `configs.py` | gripper_link local; up-forward-back |
169+
170+
## CLI usage
171+
172+
```bash
173+
# Smoke-test the scene + expert; writes 3 MP4s (agentview, wrist, composite)
174+
uv run python -m sim.scripts.record_rollout --out videos/demo.mp4 --n-episodes 3 --seed 0
175+
176+
# Collect a demo dataset
177+
uv run python -m sim.scripts.collect_demos --repo-id local/safe-cube-bc \
178+
--root data/safe_cube_bc --n-episodes 50 --successes-only
179+
180+
# Interactive viewer (macOS — note mjpython, not python)
181+
.venv/bin/mjpython -m sim.scripts.view_env
182+
```
183+
184+
## What's NOT done yet (for the next agent)
185+
186+
- **`SafePI0Policy`** — the LeRobot subclass with `L = L_flow_matching + λ·L_safety` per `project_summary.md` §5. Privileged keys are already in the recorded datasets (`privileged.cube_positions`, `privileged.cube_half_extents`, `privileged.blue_cube_pos`, `privileged.goal_pos`, `privileged.ee_pos`, `privileged.grasped`) — the policy/loss just needs to read them from the batch and ignore them at inference. Per-cube arrays are flattened to 1D; reshape on load.
187+
- **DAgger driver** — sketched in `sim/README.md`. The primitives it needs are already in place (`SafeCubeEnv`, `ScriptedExpert`, `EpisodeRecorder`, `evaluate.evaluate()`).
188+
- **Domain randomization not yet applied.** `sim/randomize.py` defines `apply_dr(model, rng, cfg)` but `env.reset()` doesn't call it. Wire it in after the scene compiles, before the first physics step, when ready.
189+
- **Real-world fine-tune / hardware deployment** — out of scope for the sim.
190+
191+
## Gotchas for future edits
192+
193+
- If you add a new visual-only geom to the world, add its name to `_NON_ARM_WORLD_GEOMS` in `scene.py` so the collision-mask rewrite doesn't flag it as a flattened-base arm geom.
194+
- If you change `grip_open` or `grip_closed`, double-check the four thresholds: `_grip_qpos_attach_threshold` / `_grip_qpos_release_threshold` (env.py), and `grasp_close_qpos_threshold` / `grasp_open_qpos_threshold` (configs.py). They need to sit between the new closed and open values with hysteresis.
195+
- After `expert.done()`, the rollout scripts must send `grip_open` (not `home_gripper`) as the gripper action — otherwise the gripper closes back on the released cube and re-attaches it via the magnetic grip. Already wired in `record_rollout.py` and `collect_demos.py`; replicate the pattern if you write new drivers.
196+
- `success` only fires after `success_dwell_steps = 5` frames of the cube resting in the goal zone *with the gripper released*. Don't exit the env loop the instant `expert.done()` is True — keep stepping for ~1 s with a held-open pose so the dwell window can register.
197+
- The wrist camera is mounted on `gripper_link` (rotates with the gripper's `wrist_roll`). If you ever change the home pose so the wrist is rolled, expect the wrist-cam horizon to tilt accordingly.
198+
- Sample-time path BFS uses `path_clearance_radius = 0.045 m`. If you shrink `min_cube_separation` or pack the field tighter, this BFS will start rejecting layouts and `sample_layout` will raise after 100 attempts.

0 commit comments

Comments
 (0)