Scheduled Commit

scripts/eval/trex_ablation_eval.py

@@ -47,6 +47,7 @@ from groot.vla.model.trex_track_force.force import (
 )
 from groot.vla.model.trex_track_force.runtime import TrexRuntimeStatistics
 from groot.vla.model.trex_track_force.track import TRACK_HORIZON
+from groot.vla.model.n1_5.sim_policy import unsqueeze_dict_values
 from groot.vla.model.trex_track_force.vla import TrexTrackForceVLA
 
 ACTION_RATE_HZ = 20.0
@@ -496,10 +497,27 @@ def main() -> None:
         for anchor_index, anchor in enumerate(openloop_anchor_times):
             raw = builder.build(anchor, blocks=1, prompt=prompt,
                                 with_deform=use_deform, history_only_video=True)
-            sample = transform_eval(dict(raw))
-            gt = np.asarray(sample["action"], dtype=np.float32)[..., :62]
-            gt = torch.as_tensor(gt.reshape(-1, 62)[:ACTION_HORIZON])
-            batch = to_batch(sample, collator, device, torch.bfloat16)
+            delta = np.asarray(raw.pop("action.eef62"), dtype=np.float32)
+            scale = builder.stats.action_q99 - builder.stats.action_q01
+            gt = np.clip(
+                2.0 * (delta - builder.stats.action_q01)
+                / np.where(scale == 0, 1.0, scale)
+                - 1.0,
+                -1.0,
+                1.0,
+            )
+            gt = np.where(scale == 0, delta, gt)
+            gt = torch.as_tensor(gt[:ACTION_HORIZON])
+            collated = transform_eval(unsqueeze_dict_values(dict(raw)))
+            batch = {}
+            for key, value in collated.items():
+                if torch.is_tensor(value):
+                    value = (
+                        value.to(device=device, dtype=torch.bfloat16)
+                        if value.is_floating_point()
+                        else value.to(device=device)
+                    )
+                batch[key] = value
             if use_deform and "deform_current" in batch:
                 batch["deform_current"] = batch["deform_current"][:, :1, :1]
             modes = [("cascade", True)] if use_force else []