From 7c942bf6ab66829964e4306b5ff141df081af1c9 Mon Sep 17 00:00:00 2001 From: William Chen Date: Mon, 15 Jun 2026 23:31:10 +0000 Subject: [PATCH] update output objects to support batching --- src/openpi/policies/droid_policy.py | 2 +- src/openpi/policies/libero_policy.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/openpi/policies/droid_policy.py b/src/openpi/policies/droid_policy.py index 666b7f76fe..55bdb419dd 100644 --- a/src/openpi/policies/droid_policy.py +++ b/src/openpi/policies/droid_policy.py @@ -78,4 +78,4 @@ def __call__(self, data: dict) -> dict: class DroidOutputs(transforms.DataTransformFn): def __call__(self, data: dict) -> dict: # Only return the first 8 dims. - return {"actions": np.asarray(data["actions"][:, :8])} + return {"actions": np.asarray(data["actions"][..., :8])} diff --git a/src/openpi/policies/libero_policy.py b/src/openpi/policies/libero_policy.py index 10611f61be..fe5aab0add 100644 --- a/src/openpi/policies/libero_policy.py +++ b/src/openpi/policies/libero_policy.py @@ -97,4 +97,4 @@ def __call__(self, data: dict) -> dict: # dimension, we need to now parse out the correct number of actions in the return dict. # For Libero, we only return the first 7 actions (since the rest is padding). # For your own dataset, replace `7` with the action dimension of your dataset. - return {"actions": np.asarray(data["actions"][:, :7])} + return {"actions": np.asarray(data["actions"][..., :7])}