Fl

.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

@@ -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,18 @@ 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.
+
 ### 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,11 @@ 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. |
 | `--resume` | Append to an existing dataset (`EXPORT_SUCCEEDED_ONLY_RESUME`) instead of starting fresh. |
 | `--seed` | Optional. Defaults to `int(time.time())`. |
 
@@ -182,7 +189,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 +227,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/utils/object_pose_augmentation.py

@@ -0,0 +1,178 @@
+"""Pose augmentation utilities for synthetic data generation.
+
+These helpers operate on the world-frame episode poses returned by
+``load_episode_poses``. The goal is to expand a small number of UMI-derived
+scene setups into more simulator rollouts without depending on the raw UMI
+trajectory, which is not used by the current training pipeline.
+"""
+
+from __future__ import annotations
+
+import math
+import random
+
+from simulator.utils.object_poses_loader import WorldPose
+
+EpisodeWorldPoses = dict[str, WorldPose]
+
+
+class PoseAugmentationError(ValueError):
+    """Raised when pose augmentation inputs are malformed."""
+
+
+def augment_episode_world_poses(
+    episodes: list[EpisodeWorldPoses],
+    *,
+    factor: int,
+    seed: int,
+    global_xy_jitter: float = 0.0,
+    local_xy_jitter: float = 0.0,
+    yaw_jitter_deg: float = 0.0,
+    min_object_distance: float = 0.0,
+    mix_objects: bool = False,
+    max_attempts: int = 64,
+) -> list[EpisodeWorldPoses]:
+    """Expand replay episodes with conservative world-frame perturbations.
+
+    Args:
+        episodes: Base replay episodes from ``load_episode_poses``.
+        factor: Dataset multiplier. ``1`` returns a copy of the original list.
+        seed: RNG seed for deterministic augmentation.
+        global_xy_jitter: Uniform translation range applied to the full scene.
+        local_xy_jitter: Additional per-object uniform translation range.
+        yaw_jitter_deg: Uniform world-yaw jitter applied to every object.
+        min_object_distance: Minimum allowed XY distance between any two objects.
+        mix_objects: If True, build new episodes by sampling each object from the
+            full per-object pose bank instead of preserving original pairings.
+        max_attempts: Rejection-sampling attempts per synthetic episode.
+    """
+    if factor < 1:
+        raise PoseAugmentationError(f"factor must be >= 1, got {factor}")
+    if max_attempts < 1:
+        raise PoseAugmentationError(f"max_attempts must be >= 1, got {max_attempts}")
+    if any(value < 0.0 for value in (global_xy_jitter, local_xy_jitter, yaw_jitter_deg, min_object_distance)):
+        raise PoseAugmentationError("jitter and min_object_distance values must be non-negative")
+    if not episodes:
+        return []
+
+    object_names = tuple(sorted(episodes[0].keys()))
+    if not object_names:
+        raise PoseAugmentationError("episodes must contain at least one object")
+    for ep_idx, episode in enumerate(episodes[1:], start=1):
+        names = tuple(sorted(episode.keys()))
+        if names != object_names:
+            raise PoseAugmentationError(
+                f"episode {ep_idx} object set {names} does not match episode 0 object set {object_names}"
+            )
+
+    rng = random.Random(seed)
+    pose_bank = {
+        name: [_clone_world_pose(episode[name]) for episode in episodes]
+        for name in object_names
+    }
+    out = [_clone_episode(episode) for episode in episodes]
+
+    if factor == 1:
+        return out
+
+    yaw_jitter_rad = math.radians(yaw_jitter_deg)
+    for episode in episodes:
+        for _ in range(factor - 1):
+            out.append(
+                _sample_augmented_episode(
+                    base_episode=episode,
+                    pose_bank=pose_bank,
+                    object_names=object_names,
+                    rng=rng,
+                    global_xy_jitter=global_xy_jitter,
+                    local_xy_jitter=local_xy_jitter,
+                    yaw_jitter_rad=yaw_jitter_rad,
+                    min_object_distance=min_object_distance,
+                    mix_objects=mix_objects,
+                    max_attempts=max_attempts,
+                )
+            )
+    return out
+
+
+def _sample_augmented_episode(
+    *,
+    base_episode: EpisodeWorldPoses,
+    pose_bank: dict[str, list[WorldPose]],
+    object_names: tuple[str, ...],
+    rng: random.Random,
+    global_xy_jitter: float,
+    local_xy_jitter: float,
+    yaw_jitter_rad: float,
+    min_object_distance: float,
+    mix_objects: bool,
+    max_attempts: int,
+) -> EpisodeWorldPoses:
+    for _ in range(max_attempts):
+        dx_global = rng.uniform(-global_xy_jitter, global_xy_jitter)
+        dy_global = rng.uniform(-global_xy_jitter, global_xy_jitter)
+        yaw_delta = rng.uniform(-yaw_jitter_rad, yaw_jitter_rad) if yaw_jitter_rad > 0.0 else 0.0
+        episode: EpisodeWorldPoses = {}
+
+        for name in object_names:
+            source_pose = rng.choice(pose_bank[name]) if mix_objects else base_episode[name]
+            pos, quat = source_pose
+            dx_local = rng.uniform(-local_xy_jitter, local_xy_jitter)
+            dy_local = rng.uniform(-local_xy_jitter, local_xy_jitter)
+            episode[name] = (
+                (pos[0] + dx_global + dx_local, pos[1] + dy_global + dy_local, pos[2]),
+                _rotate_world_yaw(quat, yaw_delta),
+            )
+
+        if _has_valid_separation(episode, min_object_distance):
+            return episode
+
+    return _clone_episode(base_episode)
+
+
+def _has_valid_separation(episode: EpisodeWorldPoses, min_object_distance: float) -> bool:
+    if min_object_distance <= 0.0 or len(episode) < 2:
+        return True
+    poses = list(episode.values())
+    min_dist_sq = min_object_distance * min_object_distance
+    for idx, (pos_a, _) in enumerate(poses):
+        for pos_b, _ in poses[idx + 1 :]:
+            dx = pos_a[0] - pos_b[0]
+            dy = pos_a[1] - pos_b[1]
+            if dx * dx + dy * dy < min_dist_sq:
+                return False
+    return True
+
+
+def _rotate_world_yaw(quat_wxyz: tuple[float, float, float, float], yaw_delta: float) -> tuple[float, float, float, float]:
+    if abs(yaw_delta) < 1e-12:
+        return tuple(float(v) for v in quat_wxyz)
+    half = yaw_delta * 0.5
+    delta = (math.cos(half), 0.0, 0.0, math.sin(half))
+    return _quat_mul_wxyz(delta, quat_wxyz)
+
+
+def _quat_mul_wxyz(
+    q1: tuple[float, float, float, float],
+    q2: tuple[float, float, float, float],
+) -> tuple[float, float, float, float]:
+    w1, x1, y1, z1 = q1
+    w2, x2, y2, z2 = q2
+    return (
+        w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
+        w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
+        w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
+        w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
+    )
+
+
+def _clone_episode(episode: EpisodeWorldPoses) -> EpisodeWorldPoses:
+    return {name: _clone_world_pose(pose) for name, pose in episode.items()}
+
+
+def _clone_world_pose(world_pose: WorldPose) -> WorldPose:
+    pos, quat = world_pose
+    return (
+        (float(pos[0]), float(pos[1]), float(pos[2])),
+        (float(quat[0]), float(quat[1]), float(quat[2]), float(quat[3])),
+    )

run.sh

@@ -0,0 +1,67 @@
+make launch-isaaclab-glowsai-l40s
+# data generation
+python scripts/datagen/generate.py \
+    --task HCIS-CutleryArrangement-SingleArm-v0 \
+    --num_envs 1 \
+    --device cuda \
+    --headless \
+    --enable_cameras \
+    --record \
+    --use_lerobot_recorder \
+    --lerobot_dataset_repo_id fanyi000/cultery_synth \
+    --augment_pose_factor 10 \
+    --augment_global_xy_jitter 0.01 \
+    --augment_local_xy_jitter 0.05 \
+    --object_poses data/AI-final-49/object_poses.json
+
+# use some dummy data
+python scripts/datagen/generate.py \
+    --task HCIS-CutleryArrangement-SingleArm-v0 \
+    --num_envs 1 \
+    --device cuda \
+    --headless \
+    --enable_cameras \
+    --record \
+    --use_lerobot_recorder \
+    --lerobot_dataset_repo_id fanyi000/cultery_synth \
+    --augment_pose_factor 10 \
+    --augment_global_xy_jitter 0.01 \
+    --augment_local_xy_jitter 0.05 \
+    --object_poses data/AI-final-49/object_poses_combined.json
+
+# training
+uv sync && source .venv/bin/activate
+HF_HUB_DISABLE_XET=1 lerobot-train \
+  --dataset.repo_id=${HF_USER}/cultery_synth \
+  --policy.type=diffusion \
+  --output_dir=outputs/diffusion_v2 \
+  --job_name=cupstacking \
+  --policy.device=cuda \
+  --wandb.enable=true \
+  --policy.repo_id=${HF_USER}/my_policy_diffusion
+
+# training exp1
+uv sync && source .venv/bin/activate
+HF_HUB_DISABLE_XET=1 lerobot-train \
+  --dataset.repo_id=${HF_USER}/cultery_synth \
+  --policy.type=act \
+  --output_dir=outputs/act_v1 \
+  --job_name=cupstacking \
+  --policy.device=cuda \
+  --wandb.enable=true \
+  --policy.repo_id=${HF_USER}/my_policy_act
+
+# evaluation
+hf download fanyi000/my_policy_diffusion --local-dir outputs/diffusion_v3
+make launch-isaaclab-glowsai-l40s
+python scripts/rollout.py \
+    --task=eval/cutlery_arrangement_eval.py \
+    --policy_type=lerobot-diffusion \
+    --policy_checkpoint_path=outputs/diffusion_v3 \
+    --policy_action_horizon=8 \
+    --device=cuda \
+    --headless \
+    --enable_cameras \
+    --eval_rounds=30 \
+    --episode_length_s=60 \
+    2>&1 | tee /workspace/aicapstone/eval_results.txt