Upload tinyvla2/scripts/eval_libero_canon.py with huggingface_hub

tinyvla2/scripts/eval_libero_canon.py

@@ -47,6 +47,14 @@ def main():
     p.add_argument("--n-action-steps", type=int, default=25,
                    help="actions executed per predicted chunk before re-planning")
     p.add_argument("--grip-mode", default="threshold", choices=["threshold", "linear"])
+    p.add_argument("--decode", default="servo", choices=["servo", "scaled"],
+                   help="servo (default): integrate predicted deltas into a desired EE pose and "
+                        "command the error against the measured pose. The naive scaled decode "
+                        "replays achieved motion as a command, which the OSC controller "
+                        "under-tracks: GT round-trip replay scores 1/5 scaled vs 3/5 servo "
+                        "against a 4/5 raw-GT ceiling.")
+    p.add_argument("--kp", type=float, default=120.0)
+    p.add_argument("--kr", type=float, default=12.0)
     p.add_argument("--save-video-dir", type=Path, default=None)
     args = p.parse_args()
 
@@ -117,18 +125,30 @@ def main():
                 "morphology": morph,
                 "embodiment_id": torch.tensor([0], device="cuda")}
 
-    def decode(chunk_norm):
-        """(chunk, >=7) normalized canonical -> (chunk-1, 7) OSC actions."""
-        phys = chunk_norm[:, :7] * span / 2 + mid
-        phys = phys[1:]  # drop the identically-zero anchor delta
+    def to_physical(chunk_norm):
+        """(chunk, >=7) normalized canonical -> (chunk-1, 7) physical deltas.
+
+        Drops the anchor delta: sidecars built from observation state have
+        a_pose == s_pose, so chunk[0] is identically zero.
+        """
+        return (chunk_norm[:, :7] * span / 2 + mid)[1:]
+
+    def grip_cmd(openness):
+        if args.grip_mode == "threshold":
+            return -1.0 if openness > fit["grip_threshold"] else 1.0
+        return float(np.clip(openness * fit["grip_a"] + fit["grip_b"], -1, 1))
+
+    def decode_scaled(phys):
         osc = np.zeros((len(phys), 7), dtype=np.float32)
         osc[:, :6] = phys[:, :6] * mscale
-        if args.grip_mode == "threshold":
-            osc[:, 6] = np.where(phys[:, 6] > fit["grip_threshold"], -1.0, 1.0)
-        else:
-            osc[:, 6] = phys[:, 6] * fit["grip_a"] + fit["grip_b"]
+        osc[:, 6] = [grip_cmd(g) for g in phys[:, 6]]
         return np.clip(osc, -1.0, 1.0)
 
+    def measured_pose(obs):
+        rs = obs["robot_state"]
+        return (np.asarray(rs["eef"]["pos"]).flatten(),
+                Rotation.from_quat(np.asarray(rs["eef"]["quat"]).flatten()))
+
     fps = meta.fps
     total = succ = 0
     per_task = {}
@@ -149,6 +169,7 @@ def main():
         frames = [] if args.save_video_dir else None
         latent, last_refresh = None, -1e9
         queue, done, step_i, ep_succ = [], False, 0, False
+        des_p, des_R = None, None
         while not done and step_i < args.max_steps:
             if not queue:
                 batch = to_batch(obs, task_text)
@@ -159,8 +180,24 @@ def main():
                     batch["semantic_latent"] = latent
                 with torch.autocast("cuda", torch.bfloat16):
                     pred = policy.predict_action_chunk(batch)[0].float().cpu().numpy()
-                queue = list(decode(pred)[: args.n_action_steps])
-            act = queue.pop(0)
+                phys = to_physical(pred)[: args.n_action_steps]
+                if args.decode == "servo":
+                    queue = list(phys)          # consumed as deltas, servoed below
+                    des_p, des_R = measured_pose(obs)   # re-anchor each re-plan
+                else:
+                    queue = list(decode_scaled(phys))
+            if args.decode == "servo":
+                d = queue.pop(0)
+                cur_p, cur_R = measured_pose(obs)
+                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] = grip_cmd(d[6])
+                act = np.clip(act, -1, 1)
+            else:
+                act = queue.pop(0)
             try:
                 obs, _, terminated, truncated, info = env.step(act[None])
             except ValueError: