Refine dataset filtering and ground alignment

Author: BotRunner64

teleopit/inputs/pico4_provider.py

@@ -121,15 +121,15 @@ def _has_non_degenerate_positions(positions: NDArray[np.float64]) -> bool:
     return extent > 1e-6
 
 
-def _compute_ground_lift_offset(positions: NDArray[np.float64]) -> float:
+def _compute_ground_alignment_offset(positions: NDArray[np.float64]) -> float:
     pos = np.asarray(positions, dtype=np.float64).reshape(-1, 3)
     if pos.size == 0:
         return 0.0
     finite_mask = np.all(np.isfinite(pos), axis=1)
     if not np.any(finite_mask):
         return 0.0
     min_z = float(np.min(pos[finite_mask, 2]))
-    return max(-min_z, 0.0)
+    return -min_z
 
 
 def _bridge_accepts_video_enabled(bridge_cls: type[Any]) -> bool:
@@ -233,7 +233,7 @@ def __init__(
         self._last_source_seq: int | None = None
         self._controller_snapshot: PicoControllerSnapshot | None = None
         self._hand_snapshot: PicoHandSnapshot | None = None
-        self._ground_lift_offset: float | None = None
+        self._ground_alignment_offset: float | None = None
         self._bridge = bridge_cls(
             host=bridge_host,
             port=int(bridge_port),
@@ -415,15 +415,15 @@ def _accept_pico_frame(self, frame: Any) -> bool:
                 and timestamp - self._last_frame_timestamp > self._timestamp_gap_reset_s
             ):
                 self._frame_cache.clear()
-                self._ground_lift_offset = None
+                self._ground_alignment_offset = None
                 logger.warning(
                     "Pico4InputProvider timestamp-gap reset | gap=%.4fs",
                     timestamp - self._last_frame_timestamp,
                 )
             if self._last_frame_timestamp is not None and timestamp <= self._last_frame_timestamp + 1e-9:
                 timestamp = self._last_frame_timestamp + 1e-6
 
-            human_frame = self._apply_ground_lift(human_frame)
+            human_frame = self._apply_ground_alignment(human_frame)
             self._frame_cache.append(human_frame, timestamp, fps_timestamp=timestamp)
             self._last_raw_body_joints = body_joints.copy()
             self._last_frame_timestamp = timestamp
@@ -534,17 +534,17 @@ def _convert_body_joints_to_frame(body_joints: NDArray[np.float64]) -> HumanFram
             result[name] = (np.asarray(pos, dtype=np.float64), np.asarray(quat, dtype=np.float64))
         return result
 
-    def _apply_ground_lift(self, human_frame: HumanFrame) -> HumanFrame:
-        """Apply one fixed Z lift so the initial Pico skeleton sits on the floor."""
-        if self._ground_lift_offset is None:
+    def _apply_ground_alignment(self, human_frame: HumanFrame) -> HumanFrame:
+        """Apply one fixed Z offset so the initial Pico skeleton sits on the floor."""
+        if self._ground_alignment_offset is None:
             positions = np.asarray([value[0] for value in human_frame.values()], dtype=np.float64)
             if _has_non_degenerate_positions(positions):
-                self._ground_lift_offset = _compute_ground_lift_offset(positions)
+                self._ground_alignment_offset = _compute_ground_alignment_offset(positions)
             else:
                 return human_frame
 
-        offset = float(self._ground_lift_offset)
-        if offset <= 0.0:
+        offset = float(self._ground_alignment_offset)
+        if abs(offset) <= 1e-12:
             return human_frame
 
         z_offset = np.array([0.0, 0.0, offset], dtype=np.float64)

teleopit/runtime/assets.py

@@ -6,7 +6,6 @@
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
 GMR_ASSETS_ROOT = PROJECT_ROOT / "teleopit" / "retargeting" / "gmr" / "assets"
 UNITREE_G1_MJLAB_XML = GMR_ASSETS_ROOT / "unitree_g1" / "g1_mjlab.xml"
-UNITREE_G1_MJLAB_PAYLOAD_XML = GMR_ASSETS_ROOT / "unitree_g1" / "g1_mjlab_payload.xml"
 
 
 def missing_gmr_assets_message(path: str | Path, *, label: str = "Required asset") -> str:

tests/test_dataset_v2.py

@@ -131,20 +131,26 @@ def test_load_dataset_spec_parses_preprocess(tmp_path: Path) -> None:
 hash_salt: ""
 preprocess:
   normalize_root_xy: true
-  ground_align: clip_min_foot
+  ground_align: none
   min_frames: 10
+  max_all_off_ground_s: 0.5
+  off_ground_height: 0.08
 sources:
   - name: clips
     type: npz
     input: {tmp_path / 'npz_source'}
+    exclude_patterns: ["*obstacle*"]
 """,
         encoding="utf-8",
     )
 
     spec = load_dataset_spec(spec_path)
     assert spec.preprocess.normalize_root_xy is True
-    assert spec.preprocess.ground_align == "clip_min_foot"
+    assert spec.preprocess.ground_align == "none"
     assert spec.preprocess.min_frames == 10
+    assert spec.preprocess.max_all_off_ground_s == 0.5
+    assert spec.preprocess.off_ground_height == 0.08
+    assert spec.sources[0].exclude_patterns == ("*obstacle*",)
 
 
 def test_load_dataset_spec_parses_seed_filter_preset(tmp_path: Path) -> None:
@@ -428,6 +434,71 @@ def test_collect_source_files_with_report_handles_single_file_source(tmp_path: P
     assert [item.rel_no_suffix.as_posix() for item in legacy_items] == ["clip_a"]
 
 
+def test_collect_source_files_with_report_applies_exclude_patterns(tmp_path: Path) -> None:
+    input_root = tmp_path / "lafan1"
+    input_root.mkdir(parents=True, exist_ok=True)
+    (input_root / "walk1.bvh").write_text("placeholder", encoding="utf-8")
+    (input_root / "obstacle_run.bvh").write_text("placeholder", encoding="utf-8")
+    nested = input_root / "subject1"
+    nested.mkdir()
+    (nested / "obstacle_jump.bvh").write_text("placeholder", encoding="utf-8")
+
+    source = DatasetSourceSpec(
+        name="lafan1",
+        type="bvh",
+        input=str(input_root),
+        bvh_format="lafan1",
+        exclude_patterns=("*obstacle*",),
+    )
+
+    items, _scan_root, report = dataset_builder._collect_source_files_with_report(
+        source,
+        quiet=True,
+    )
+
+    assert [item.rel_no_suffix.as_posix() for item in items] == ["walk1"]
+    assert report["scanned_files"] == 3
+    assert report["path_rejected_files"] == 2
+    assert report["kept_files"] == 1
+    assert report["filtered_files"] == 2
+    assert report["path_reject_reasons"] == {"*obstacle*": 2}
+
+
+def test_collect_source_files_with_report_preserves_path_excludes_with_metadata(tmp_path: Path) -> None:
+    input_root = tmp_path / "seed_source" / "g1" / "csv"
+    input_root.mkdir(parents=True, exist_ok=True)
+    for name in ("walk_a.csv", "walk_b.csv", "obstacle_walk.csv"):
+        (input_root / name).write_text("placeholder", encoding="utf-8")
+
+    metadata_csv = tmp_path / "metadata.csv"
+    with metadata_csv.open("w", encoding="utf-8", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=["move_g1_path", "is_mirror"])
+        writer.writeheader()
+        for name in ("walk_a.csv", "walk_b.csv", "obstacle_walk.csv"):
+            writer.writerow({"move_g1_path": f"g1/csv/{name}", "is_mirror": "False"})
+
+    source = DatasetSourceSpec(
+        name="seed",
+        type="seed_csv",
+        input=str(input_root),
+        metadata_csv=str(metadata_csv),
+        filters={"is_mirror": [False]},
+        exclude_patterns=("*obstacle*",),
+    )
+
+    items, _scan_root, report = dataset_builder._collect_source_files_with_report(
+        source,
+        quiet=True,
+    )
+
+    assert [item.rel_no_suffix.as_posix() for item in items] == ["walk_a", "walk_b"]
+    assert report["scanned_files"] == 3
+    assert report["path_rejected_files"] == 1
+    assert report["kept_files"] == 2
+    assert report["filtered_files"] == 1
+    assert report["path_reject_reasons"] == {"*obstacle*": 1}
+
+
 def test_build_dataset_from_spec_writes_shard_directories(tmp_path: Path) -> None:
     npz_input = tmp_path / "npz_source"
     _write_npz_from_pkl(npz_input / "clip_a.npz")
@@ -460,6 +531,35 @@ def test_build_dataset_from_spec_writes_shard_directories(tmp_path: Path) -> Non
     assert "clip_lengths" in train_data.files
 
 
+def test_collect_clip_rows_ignores_stale_excluded_cached_npz(tmp_path: Path) -> None:
+    npz_input = tmp_path / "npz_source"
+    for name in ("keep_a.npz", "keep_b.npz", "obstacle_old.npz"):
+        _write_npz_from_pkl(npz_input / name)
+
+    spec = DatasetSpec(
+        name="demo_dataset",
+        target_fps=30,
+        val_percent=5,
+        hash_salt="",
+        sources=[
+            DatasetSourceSpec(
+                name="npz_src",
+                type="npz",
+                input=str(npz_input),
+                exclude_patterns=("*obstacle*",),
+            )
+        ],
+    )
+    paths = dataset_builder.resolve_dataset_paths(spec, output_root=tmp_path / "datasets")
+    source_dir = paths.clips_root / "npz_src"
+    for name in ("keep_a.npz", "keep_b.npz", "obstacle_old.npz"):
+        _write_npz_from_pkl(source_dir / name)
+
+    rows = dataset_builder.collect_clip_rows(spec, paths=paths)
+
+    assert sorted(row.clip_id for row in rows) == ["npz_src:keep_a", "npz_src:keep_b"]
+
+
 def test_convert_source_to_npz_clips_applies_preprocess(tmp_path: Path) -> None:
     npz_input = tmp_path / "npz_source"
     _write_npz_from_pkl(npz_input / "clip_a.npz")
@@ -503,6 +603,66 @@ def test_convert_source_to_npz_clips_applies_preprocess(tmp_path: Path) -> None:
     assert np.isclose(float(np.min(foot_z)), 0.0)
 
 
+def test_convert_source_to_npz_clips_skips_all_off_ground_clips_before_ground_align(tmp_path: Path) -> None:
+    npz_input = tmp_path / "npz_source"
+    _write_npz_from_pkl(npz_input / "keep.npz")
+    _write_npz_from_pkl(npz_input / "float.npz")
+
+    for name, floating in (("keep.npz", False), ("float.npz", True)):
+        path = npz_input / name
+        clip = dict(np.load(path, allow_pickle=True))
+        body_names = [str(body_name) for body_name in clip["body_names"].tolist()]
+        left_idx = body_names.index("left_ankle_roll_link")
+        right_idx = body_names.index("right_ankle_roll_link")
+        body_pos_w = np.asarray(clip["body_pos_w"]).copy()
+        if floating:
+            body_pos_w[..., 2] = np.maximum(body_pos_w[..., 2], 0.3)
+        else:
+            body_pos_w[:, left_idx, 2] = 0.0
+            body_pos_w[:, right_idx, 2] = 0.3
+        clip["body_pos_w"] = body_pos_w
+        np.savez(path, **clip)
+
+    source = DatasetSourceSpec(name="npz_src", type="npz", input=str(npz_input))
+    output_dir = tmp_path / "dataset" / "clips" / "npz_src"
+    report = convert_source_to_npz_clips(
+        source,
+        output_dir,
+        jobs=1,
+        preprocess=dataset_builder.DatasetPreprocessSpec(
+            ground_align="clip_min_foot",
+            max_all_off_ground_s=0.05,
+            off_ground_height=0.08,
+        ),
+    )
+
+    assert report["clips"] == 1
+    assert (output_dir / "keep.npz").is_file()
+    assert not (output_dir / "float.npz").exists()
+    assert (output_dir / "float.npz.filtered.json").is_file()
+
+    def _unexpected_run_conversion_tasks(*_args, **_kwargs):
+        raise AssertionError("filtered clip should be skipped by marker on incremental rebuild")
+
+    monkeypatch = pytest.MonkeyPatch()
+    monkeypatch.setattr(dataset_builder, "run_conversion_tasks", _unexpected_run_conversion_tasks)
+    try:
+        second_report = convert_source_to_npz_clips(
+            source,
+            output_dir,
+            jobs=1,
+            preprocess=dataset_builder.DatasetPreprocessSpec(
+                ground_align="clip_min_foot",
+                max_all_off_ground_s=0.05,
+                off_ground_height=0.08,
+            ),
+        )
+    finally:
+        monkeypatch.undo()
+
+    assert second_report["clips"] == 1
+
+
 def test_merge_clip_dicts_rejects_inconsistent_body_names(tmp_path: Path) -> None:
     clip_a = tmp_path / "clip_a.npz"
     clip_b = tmp_path / "clip_b.npz"

tests/test_domain_randomization.py

@@ -19,10 +19,6 @@ def test_general_tracking_domain_randomization_matches_gr00t_active_set() -> Non
         "add_joint_default_pos",
         "physics_material",
         "randomize_rigid_body_mass",
-        "randomize_dexhand_payload_mass",
-        "randomize_gimbal_payload_mass",
-        "randomize_dexhand_payload_pos",
-        "randomize_gimbal_payload_pos",
     }
 
     push_robot = events["push_robot"]
@@ -69,46 +65,6 @@ def test_general_tracking_domain_randomization_matches_gr00t_active_set() -> Non
     assert mass.params["asset_cfg"].body_names == "torso_link"
     assert mass.params["alpha_range"] == (-0.1, 0.45)
 
-    dexhand_mass = events["randomize_dexhand_payload_mass"]
-    assert dexhand_mass.func is dr.pseudo_inertia
-    assert dexhand_mass.mode == "startup"
-    assert dexhand_mass.params["asset_cfg"].body_names == (
-        "left_dexhand_payload",
-        "right_dexhand_payload",
-    )
-    assert dexhand_mass.params["alpha_range"] == (-1, 0)
-
-    gimbal_mass = events["randomize_gimbal_payload_mass"]
-    assert gimbal_mass.func is dr.pseudo_inertia
-    assert gimbal_mass.mode == "startup"
-    assert gimbal_mass.params["asset_cfg"].body_names == ("head_gimbal_payload",)
-    assert gimbal_mass.params["alpha_range"] == (-1, 0)
-
-    dexhand_pos = events["randomize_dexhand_payload_pos"]
-    assert dexhand_pos.func is dr.body_pos
-    assert dexhand_pos.mode == "startup"
-    assert dexhand_pos.params["asset_cfg"].body_names == (
-        "left_dexhand_payload",
-        "right_dexhand_payload",
-    )
-    assert dexhand_pos.params["operation"] == "abs"
-    assert dexhand_pos.params["ranges"] == {
-        0: (0.055, 0.095),
-        1: (-0.02, 0.02),
-        2: (-0.02, 0.02),
-    }
-
-    gimbal_pos = events["randomize_gimbal_payload_pos"]
-    assert gimbal_pos.func is dr.body_pos
-    assert gimbal_pos.mode == "startup"
-    assert gimbal_pos.params["asset_cfg"].body_names == ("head_gimbal_payload",)
-    assert gimbal_pos.params["operation"] == "abs"
-    assert gimbal_pos.params["ranges"] == {
-        0: (0.05, 0.09),
-        1: (-0.02, 0.02),
-        2: (0.43, 0.47),
-    }
-
 
 def test_play_env_disables_training_only_domain_randomization() -> None:
     import mjlab.tasks  # noqa: F401
@@ -122,8 +78,4 @@ def test_play_env_disables_training_only_domain_randomization() -> None:
     assert "add_joint_default_pos" not in play_cfg.events
     assert "physics_material" not in play_cfg.events
     assert "randomize_rigid_body_mass" not in play_cfg.events
-    assert "randomize_dexhand_payload_mass" not in play_cfg.events
-    assert "randomize_gimbal_payload_mass" not in play_cfg.events
-    assert "randomize_dexhand_payload_pos" not in play_cfg.events
-    assert "randomize_gimbal_payload_pos" not in play_cfg.events
     assert play_cfg.events == {}