Upload tinyvla2/scripts/replay_libero_canon.py with huggingface_hub

tinyvla2/scripts/replay_libero_canon.py

@@ -26,11 +26,19 @@ def main():
     ap.add_argument("--episodes", type=int, default=5)
     ap.add_argument("--grip-mode", default="threshold", choices=["threshold", "linear", "gt"])
     ap.add_argument("--raw", action="store_true", help="replay raw GT actions instead (control)")
+    ap.add_argument("--decode", default="scaled", choices=["scaled", "servo"],
+                    help="scaled: per-step delta x fitted gain (open loop). "
+                         "servo: integrate deltas into a desired pose and command the "
+                         "error against the MEASURED pose (closes the controller lag)")
+    ap.add_argument("--kp", type=float, default=120.0)
+    ap.add_argument("--kr", type=float, default=12.0)
     args = ap.parse_args()
 
     from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
     from lerobot.envs.factory import make_env, make_env_config
 
+    from scipy.spatial.transform import Rotation
+
     from tinyvla.data.canonical import CanonicalChunkStore
 
     meta = LeRobotDatasetMetadata("libero_local", root=ROOT)
@@ -78,9 +86,27 @@ def main():
         tid, env = env_by_task[task]
         obs, _ = env.reset(seed=0)
         ep_succ = False
+        des_p, des_R = None, None
         for idx in range(s, e - 1):
             if args.raw:
                 act = all_actions[idx]
+            elif args.decode == "servo":
+                # integrate the canonical deltas into an absolute desired EE pose,
+                # then command the error against what the arm actually reached
+                d = store.chunk_for(ep, idx - s)[1]
+                rs = obs["robot_state"]
+                cur_p = np.asarray(rs["eef"]["pos"]).flatten()
+                cur_R = Rotation.from_quat(np.asarray(rs["eef"]["quat"]).flatten())
+                if des_p is None:
+                    des_p, des_R = cur_p.copy(), cur_R
+                des_p = des_p + d[:3]
+                des_R = Rotation.from_rotvec(d[3:6]) * des_R
+                act = np.zeros(7, dtype=np.float32)
+                act[:3] = args.kp * (des_p - cur_p)
+                act[3:6] = args.kr * (des_R * cur_R.inv()).as_rotvec()
+                act[6] = (all_actions[idx][6] if args.grip_mode == "gt"
+                          else (-1.0 if d[6] > fit["grip_threshold"] else 1.0))
+                act = np.clip(act, -1, 1)
             else:
                 # chunk[1] is the motion caused by the action at this frame
                 act = decode(store.chunk_for(ep, idx - s)[1], all_actions[idx][6])
@@ -96,7 +122,7 @@ def main():
         succ += int(ep_succ)
         print(f"ep {ep} (task {tid}): {'SUCCESS' if ep_succ else 'fail'} | {task[:55]}", flush=True)
 
-    tag = "RAW GT" if args.raw else f"decoded canonical (grip={args.grip_mode})"
+    tag = "RAW GT" if args.raw else f"decoded canonical ({args.decode}, grip={args.grip_mode})"
     print(f"\n{tag} replay success: {succ}/{replayed}")