Gazn cutlery aug

.gitignore

@@ -7,7 +7,7 @@ __pycache__/
 build/
 dist/
 *.egg-info/
-
+outputs
 # IDE
 .idea/
 .vscode/
@@ -38,3 +38,5 @@ datasets/*
 # Private external task files
 private_tasks/*
 !private_tasks/.gitkeep
+
+*.txt

Dockerfile

@@ -135,4 +135,8 @@ RUN printf "numpy==1.26.0\n" > /tmp/sim-constraints.txt \
 RUN python -m pip install --upgrade pip==26.0.1 \
     && python -m pip install --no-deps numpy==1.26.0
 
+# Fix Vulkan ICD for headless EGL rendering (no X server)
+RUN echo '{"file_format_version":"1.0.0","ICD":{"library_path":"libEGL_nvidia.so.0","api_version":"1.4.312"}}' \
+    > /etc/vulkan/icd.d/nvidia_icd.json
+
 CMD ["/bin/bash"]

Makefile

@@ -103,6 +103,7 @@ launch-isaaclab-glowsai-4090: build-isaaclab
 	docker run --rm -it \
 		--name $(CONTAINER_NAME)-glowsai-4090 \
 		--gpus '"device=0"' \
+		--device /dev/dri/card0:/dev/dri/card0 \
 		--net=host \
 		--ipc=host \
 		--ulimit memlock=-1 \
@@ -114,7 +115,6 @@ launch-isaaclab-glowsai-4090: build-isaaclab
 		-v /tmp/.X11-unix:/tmp/.X11-unix:rw \
 		-v /opt/VirtualGL:/opt/VirtualGL:ro \
 		-v /usr/share/vulkan/icd.d:/usr/share/vulkan/icd.d:ro \
-		-v /etc/vulkan/icd.d:/etc/vulkan/icd.d:ro \
 		-e DISPLAY=:1 \
 		-e USE_VNC=1 \
 		-e OMNI_KIT_ACCEPT_EULA=Y \
@@ -136,12 +136,13 @@ launch-isaaclab-glowsai-4090: build-isaaclab
 			exec /bin/bash \
 		'
 
-# ---- Launch: GlowsAI L40S (VirtualGL + VNC display :1) -----------------------
+# ---- Launch: GlowsAI L40S (VirtualGL + VNC display :2) -----------------------
 launch-isaaclab-glowsai-l40s: build-isaaclab
 	@set -e; \
 	docker run --rm -it \
 		--name $(CONTAINER_NAME)-glowsai-l40s \
 		--gpus '"device=0"' \
+		--device /dev/dri/card0:/dev/dri/card0 \
 		--net=host \
 		--ipc=host \
 		--ulimit memlock=-1 \
@@ -153,7 +154,6 @@ launch-isaaclab-glowsai-l40s: build-isaaclab
 		-v /tmp/.X11-unix:/tmp/.X11-unix:rw \
 		-v /opt/VirtualGL:/opt/VirtualGL:ro \
 		-v /usr/share/vulkan/icd.d:/usr/share/vulkan/icd.d:ro \
-		-v /etc/vulkan/icd.d:/etc/vulkan/icd.d:ro \
 		-e DISPLAY=:1 \
 		-e USE_VNC=1 \
 		-e VGL_DISPLAY=egl0 \

README.md

@@ -127,9 +127,14 @@ python scripts/datagen/generate.py \
     --record \
     --use_lerobot_recorder \
     --lerobot_dataset_repo_id ${HF_USER}/<repo_id> \
+    --augment_pose_factor 10 \
+    --augment_global_xy_jitter 0.01 \
+    --augment_local_xy_jitter 0.01 \
     --object_poses data/<demo_directory_name>/object_poses.json
 ```
 
+`--object_poses` still provides the base UMI scene setups, but `--augment_pose_factor` lets Step 3 replay more synthetic episodes from the same small set of demonstrations. For example, `16` source entries with `--augment_pose_factor 10` become `160` replay episodes.
+
 Upload the recorded dataset to Hugging Face Hub:
 
 ```bash

docs/getting_started.md

@@ -239,6 +239,8 @@ hf download ${HF_USER}/<repo_id> --local-dir data/<demo_directory_name>
 
 ### Run data generation
 
+`object_poses.json` only supplies the per-episode object placements used to initialize the scene. If your UMI pipeline only gives you a small number of usable entries, expand them here with pose augmentation before training.
+
 Available tasks:
 - `HCIS-CupStacking-SingleArm-v0`
 - `HCIS-CutleryArrangement-SingleArm-v0`
@@ -253,9 +255,19 @@ python scripts/datagen/generate.py \
     --record \
     --use_lerobot_recorder \
     --lerobot_dataset_repo_id ${HF_USER}/<repo_id> \
+    --augment_pose_factor 10 \
+    --augment_global_xy_jitter 0.01 \
+    --augment_local_xy_jitter 0.01 \
     --object_poses data/<demo_directory_name>/object_poses.json
 ```
 
+Recommended starting point when you only have about `16` usable UMI entries:
+
+- `--augment_pose_factor 10` to turn `16` base entries into about `160` replay episodes.
+- Keep `--augment_yaw_jitter_deg 0` at first for maximum stability.
+- Add `--augment_mix_objects` only if you need more diversity after confirming the scripted policy still succeeds reliably.
+- For `HCIS-CutleryArrangement-SingleArm-v0`, Step 3 also blends a subset of augmented scenes toward the eval initial-pose distribution while keeping the rest broader, so training does not collapse onto only the public eval start states.
+
 ### Upload the generated dataset
 
 ```bash

docs/synthetic_data_generation.md

@@ -154,10 +154,13 @@ python scripts/datagen/generate.py \
     --num_envs 1 \
     --device cuda \
     --enable_cameras \
-    --num_demos 50 \
     --record \
     --use_lerobot_recorder \
-    --lerobot_dataset_repo_id HF-USER/name
+    --lerobot_dataset_repo_id HF-USER/name \
+    --augment_pose_factor 10 \
+    --augment_global_xy_jitter 0.01 \
+    --augment_local_xy_jitter 0.01 \
+    --object_poses data/<demo_directory_name>/object_poses.json
 ```
 
 Key flags:
@@ -172,7 +175,12 @@ Key flags:
 | `--use_lerobot_recorder` | Swap the default `StreamingRecorderManager` for `LeRobotRecorderManager`, which writes the LeRobot dataset format on disk instead of HDF5. |
 | `--lerobot_dataset_repo_id HF-USER/name` | Passed straight into `LeRobotDatasetCfg(repo_id=..., fps=args_cli.lerobot_dataset_fps)`. Names the on-disk dataset and the eventual HF Hub repo. |
 | `--lerobot_dataset_fps` | Frame rate the dataset is written at. Default `30`. |
-| `--num_demos 50` | Stop after **50 successful** episodes. With `--use_lerobot_recorder` the recorder runs in `EXPORT_SUCCEEDED_ONLY` mode, so failed rollouts do not count toward the target. |
+| `--object_poses data/.../object_poses.json` | Base replay set from UMI. Every `status == "full"` entry becomes one source episode before augmentation. |
+| `--augment_pose_factor 10` | Multiply the replay set in Step 3. Example: `16` source entries become `160` replay episodes. |
+| `--augment_global_xy_jitter`, `--augment_local_xy_jitter` | Scene-level and per-object translation jitter in meters. |
+| `--augment_yaw_jitter_deg` | Optional world-yaw jitter. Keep it at `0` first if you want the safest scripted rollouts. |
+| `--augment_mix_objects` | Recombine object poses across episodes before jittering. Higher diversity, slightly higher risk. |
+| `--cutlery_eval_pose_fraction`, `--cutlery_eval_pose_jitter` | For the cutlery task only: rewrite part of the augmented replay set so fork/knife starts resemble the eval initial-pose distribution, while keeping the remaining episodes broader. |
 | `--resume` | Append to an existing dataset (`EXPORT_SUCCEEDED_ONLY_RESUME`) instead of starting fresh. |
 | `--seed` | Optional. Defaults to `int(time.time())`. |
 
@@ -182,7 +190,7 @@ What happens at runtime:
 2. `_configure_env_cfg(...)` flips `env_cfg.recorders.dataset_export_mode` to `EXPORT_SUCCEEDED_ONLY` and rewires the `success` termination so the recorder controls episode endings.
 3. `_replace_recorder_manager(...)` instantiates `LeRobotRecorderManager(env_cfg.recorders, LeRobotDatasetCfg(repo_id=..., fps=...), env)`.
 4. The main loop calls `sm.pre_step → sm.get_action → env.step → sm.advance` until `sm.is_episode_done`. On episode end, `sm.check_success(env)` decides whether the recorder commits the episode.
-5. Once `recorder_manager.exported_successful_episode_count >= num_demos`, the script exits cleanly and `recorder_manager.finalize()` writes the LeRobot dataset to disk.
+5. Once every replay episode from `--object_poses` (after optional augmentation) has been attempted, the script exits cleanly and `recorder_manager.finalize()` writes the LeRobot dataset to disk.
 
 The dataset lands locally first — the `repo_id` only names the directory at this stage. Upload is a separate step.
 
@@ -220,5 +228,5 @@ Paste your `repo_id` (e.g. `HF-USER/name`) into the Space to browse the dataset.
 - [ ] `packages/simulator/src/simulator/tasks/__init__.py` imports the new subpackage.
 - [ ] State machine subclasses `leisaac.datagen.state_machine.base.StateMachineBase` and implements `setup`, `pre_step`, `get_action`, `advance`, `reset`, `check_success`, `is_episode_done`.
 - [ ] `scripts/datagen/generate.py::TASK_REGISTRY` has `<task id>: (StateMachineClass, "<teleop_device>")`.
-- [ ] Run with `--record --use_lerobot_recorder --enable_cameras --num_demos N --lerobot_dataset_repo_id HF-USER/name`.
+- [ ] Run with `--record --use_lerobot_recorder --enable_cameras --object_poses ... --lerobot_dataset_repo_id HF-USER/name`.
 - [ ] `hf auth login`, then `hf upload <repo_id> --repo-type dataset`.

packages/simulator/src/simulator/datagen/state_machine/cutlery_arrangement.py

@@ -76,9 +76,9 @@
     "panda_finger_joint2": 0.04,
 }
 
-# Pick order: fork first (place on +y / left of plate), then knife (place on -y / right)
+# Pick order: knife first, then fork.
 _PICK_ORDER = (_KNIFE_NAME, _FORK_NAME)
-_PLACE_X_SIGNS = (+1.0, -1.0)  # fork → +x of plate, knife → -x of plate
+_PLACE_X_SIGNS = (+1.0, -1.0)  # knife → +x of plate, fork → -x of plate
 
 _PHASE_DURATIONS_PER_OBJECT = (180, 130, 20, 160, 170, 15, 30)
 _PHASES_PER_OBJECT = len(_PHASE_DURATIONS_PER_OBJECT)
@@ -140,8 +140,8 @@ def _find_body_index(robot, body_name: str) -> int:
 class CutleryArrangementStateMachine(StateMachineBase):
     """Scripted Franka policy for arranging cutlery around a plate.
 
-    Picks up the fork and places it on the +y (left) side of the plate,
-    then picks up the knife and places it on the -y (right) side.
+    Picks up the knife and places it on the +x (right) side of the plate,
+    then picks up the fork and places it on the -x (left) side.
 
     Each object goes through 7 phases:
 

packages/simulator/src/simulator/tasks/cutlery_arrangement/cutlery_arrangement_env_cfg.py

@@ -89,7 +89,7 @@ def cutlery_arranged(
     knife_cfg: SceneEntityCfg,
     max_dist_xy: float,
 ) -> torch.Tensor:
-    """Termination: fork on +y side of plate, knife on -y side, both within max_dist_xy."""
+    """Termination: fork on -x side of plate, knife on +x side, both within max_dist_xy."""
     plate: RigidObject = env.scene[plate_cfg.name]
     fork: RigidObject = env.scene[fork_cfg.name]
     knife: RigidObject = env.scene[knife_cfg.name]
@@ -106,8 +106,8 @@ def cutlery_arranged(
     done = torch.logical_and(done, fork_dist_xy <= max_dist_xy)
     done = torch.logical_and(done, knife_dist_xy <= max_dist_xy)
 
-    fork_on_left = fork_pos[:, 0] > plate_pos[:, 0]
-    knife_on_right = knife_pos[:, 0] < plate_pos[:, 0]
+    fork_on_left = fork_pos[:, 0] < plate_pos[:, 0]
+    knife_on_right = knife_pos[:, 0] > plate_pos[:, 0]
 
     done = torch.logical_and(done, fork_on_left)
     done = torch.logical_and(done, knife_on_right)