Add real-world support for ACT on Aloha/Aloha2 (#228)

Co-authored-by: Alexander Soare <alexander.soare159@gmail.com>

tests/data/save_policy_to_safetensors/dora_aloha_real_act_real/actions.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2fff6294b94cf42d4dd1249dcc5c3b0269d6d9c697f894e61b867d7ab81a94e4
+size 5104

tests/data/save_policy_to_safetensors/dora_aloha_real_act_real/grad_stats.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4aa23e51607604a18b70fa42edbbe1af34f119d985628fc27cc1bbb0efbc8901
+size 31688

tests/data/save_policy_to_safetensors/dora_aloha_real_act_real/output_dict.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6fd368406c93cb562a69ff11cf7adf34a4b223507dcb2b9e9b8f44ee1036988a
+size 68

tests/data/save_policy_to_safetensors/dora_aloha_real_act_real/param_stats.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5663ee79a13bb70a1604b887dd21bf89d18482287442419c6cc6c5bf0e753e99
+size 34928

tests/data/save_policy_to_safetensors/dora_aloha_real_act_real_no_state/actions.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb1a45463efd860af2ca22c16c77d55a18bd96fef080ae77978845a2f22ef716
+size 5104

tests/data/save_policy_to_safetensors/dora_aloha_real_act_real_no_state/grad_stats.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa5a43e22f01d8e2f8d19f31753608794f1edbd74aaf71660091ab80ea58dc9b
+size 30808

tests/data/save_policy_to_safetensors/dora_aloha_real_act_real_no_state/output_dict.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:97455b4360748c99905cd103473c1a52da6901d0a73ffbc51b5ea3eb250d1386
+size 68

tests/data/save_policy_to_safetensors/dora_aloha_real_act_real_no_state/param_stats.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54d1f75cf67a7b1d7a7c6865ecb9b1cc86a2f032d1890245f8996789ab6e0df6
+size 33608

tests/scripts/save_policy_to_safetensors.py

@@ -75,15 +75,16 @@ def get_policy_stats(env_name, policy_name, extra_overrides):
     # HACK: We reload a batch with no delta_timestamps as `select_action` won't expect a timestamps dimension
     dataset.delta_timestamps = None
     batch = next(iter(dataloader))
-    obs = {
-        k: batch[k]
-        for k in batch
-        if k in ["observation.image", "observation.images.top", "observation.state"]
-    }
+    obs = {}
+    for k in batch:
+        if k.startswith("observation"):
+            obs[k] = batch[k]
+
+    if "n_action_steps" in cfg.policy:
+        actions_queue = cfg.policy.n_action_steps
+    else:
+        actions_queue = cfg.policy.n_action_repeats
 
-    actions_queue = (
-        cfg.policy.n_action_steps if "n_action_steps" in cfg.policy else cfg.policy.n_action_repeats
-    )
     actions = {str(i): policy.select_action(obs).contiguous() for i in range(actions_queue)}
     return output_dict, grad_stats, param_stats, actions
 
@@ -114,6 +115,8 @@ if __name__ == "__main__":
             ["policy.n_action_steps=8", "policy.num_inference_steps=10", "policy.down_dims=[128, 256, 512]"],
         ),
         ("aloha", "act", ["policy.n_action_steps=10"]),
+        ("dora_aloha_real", "act_real", ["policy.n_action_steps=10"]),
+        ("dora_aloha_real", "act_real_no_state", ["policy.n_action_steps=10"]),
     ]
     for env, policy, extra_overrides in env_policies:
         save_policy_to_safetensors("tests/data/save_policy_to_safetensors", env, policy, extra_overrides)

tests/test_policies.py

@@ -30,7 +30,7 @@ from lerobot.common.policies.factory import get_policy_and_config_classes, make_
 from lerobot.common.policies.normalize import Normalize, Unnormalize
 from lerobot.common.policies.policy_protocol import Policy
 from lerobot.common.utils.utils import init_hydra_config
-from tests.scripts.save_policy_to_safetensor import get_policy_stats
+from tests.scripts.save_policy_to_safetensors import get_policy_stats
 from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, require_cpu, require_env, require_x86_64_kernel
 
 
@@ -72,6 +72,8 @@ def test_get_policy_and_config_classes(policy_name: str):
         ),
         # Note: these parameters also need custom logic in the test function for overriding the Hydra config.
         ("pusht", "act", ["env.task=PushT-v0", "dataset_repo_id=lerobot/pusht"]),
+        ("dora_aloha_real", "act_real", []),
+        ("dora_aloha_real", "act_real_no_state", []),
     ],
 )
 @require_env
@@ -84,6 +86,9 @@ def test_policy(env_name, policy_name, extra_overrides):
         - Updating the policy.
         - Using the policy to select actions at inference time.
         - Test the action can be applied to the policy
+
+    Note: We test various combinations of policy and dataset. The combinations are by no means exhaustive,
+          and for now we add tests as we see fit.
     """
     cfg = init_hydra_config(
         DEFAULT_CONFIG_PATH,
@@ -135,7 +140,7 @@ def test_policy(env_name, policy_name, extra_overrides):
 
     dataloader = torch.utils.data.DataLoader(
         dataset,
-        num_workers=4,
+        num_workers=0,
         batch_size=2,
         shuffle=True,
         pin_memory=DEVICE != "cpu",
@@ -291,6 +296,8 @@ def test_normalize(insert_temporal_dim):
             ["policy.n_action_steps=8", "policy.num_inference_steps=10", "policy.down_dims=[128, 256, 512]"],
         ),
         ("aloha", "act", ["policy.n_action_steps=10"]),
+        ("dora_aloha_real", "act_real", ["policy.n_action_steps=10"]),
+        ("dora_aloha_real", "act_real_no_state", ["policy.n_action_steps=10"]),
     ],
 )
 # As artifacts have been generated on an x86_64 kernel, this test won't