WIP

tests/scripts/save_policy_to_safetensors.py

@@ -13,6 +13,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+"""
+Save the policy tests artifacts.
+
+Note: Run on the cluster
+
+Example of usage:
+```bash
+DATA_DIR=tests/data python tests/scripts/save_policy_to_safetensors.py
+```
+"""
+
+import platform
 import shutil
 from pathlib import Path
 
@@ -54,7 +66,7 @@ def get_policy_stats(env_name, policy_name, extra_overrides):
     output_dict = {k: v for k, v in output_dict.items() if isinstance(v, torch.Tensor)}
     loss = output_dict["loss"]
 
-    loss.backward()
+    loss.mean().backward()
     grad_stats = {}
     for key, param in policy.named_parameters():
         if param.requires_grad:
@@ -96,10 +108,21 @@ def save_policy_to_safetensors(output_dir, env_name, policy_name, extra_override
         print(f"Overwrite existing safetensors in '{env_policy_dir}':")
         print(f" - Validate with: `git add {env_policy_dir}`")
         print(f" - Revert with: `git checkout -- {env_policy_dir}`")
+
+    output_dict, grad_stats, param_stats, actions = get_policy_stats(env_name, policy_name, extra_overrides)
+
+    from safetensors.torch import load_file
+
+    if (env_policy_dir / "output_dict.safetensors").exists():
+        prev_loss = load_file(env_policy_dir / "output_dict.safetensors")["loss"]
+        print(f"Previous loss={prev_loss}")
+        print(f"New loss={output_dict['loss'].mean()}")
+        print()
+
+    if env_policy_dir.exists():
         shutil.rmtree(env_policy_dir)
 
     env_policy_dir.mkdir(parents=True, exist_ok=True)
-    output_dict, grad_stats, param_stats, actions = get_policy_stats(env_name, policy_name, extra_overrides)
     save_file(output_dict, env_policy_dir / "output_dict.safetensors")
     save_file(grad_stats, env_policy_dir / "grad_stats.safetensors")
     save_file(param_stats, env_policy_dir / "param_stats.safetensors")
@@ -107,27 +130,32 @@ def save_policy_to_safetensors(output_dir, env_name, policy_name, extra_override
 
 
 if __name__ == "__main__":
+    if platform.machine() != "x86_64":
+        raise OSError("Generate policy artifacts on x86_64 machine since it is used for the unit tests. ")
+
     env_policies = [
-        # ("xarm", "tdmpc", ["policy.use_mpc=false"], "use_policy"),
-        # ("xarm", "tdmpc", ["policy.use_mpc=true"], "use_mpc"),
-        # (
-        #     "pusht",
-        #     "diffusion",
-        #     [
-        #         "policy.n_action_steps=8",
-        #         "policy.num_inference_steps=10",
-        #         "policy.down_dims=[128, 256, 512]",
-        #     ],
-        #     "",
-        # ),
-        # ("aloha", "act", ["policy.n_action_steps=10"], ""),
-        # ("aloha", "act", ["policy.n_action_steps=1000", "policy.chunk_size=1000"], "_1000_steps"),
-        # ("dora_aloha_real", "act_real", ["policy.n_action_steps=10"], ""),
-        # ("dora_aloha_real", "act_real_no_state", ["policy.n_action_steps=10"], ""),
+        ("xarm", "tdmpc", ["policy.use_mpc=false"], "use_policy"),
+        ("xarm", "tdmpc", ["policy.use_mpc=true"], "use_mpc"),
+        (
+            "pusht",
+            "diffusion",
+            [
+                "policy.n_action_steps=8",
+                "policy.num_inference_steps=10",
+                "policy.down_dims=[128, 256, 512]",
+            ],
+            "",
+        ),
+        ("aloha", "act", ["policy.n_action_steps=10"], ""),
+        ("aloha", "act", ["policy.n_action_steps=1000", "policy.chunk_size=1000"], "_1000_steps"),
+        ("dora_aloha_real", "act_real", ["policy.n_action_steps=10"], ""),
+        ("dora_aloha_real", "act_real_no_state", ["policy.n_action_steps=10"], ""),
     ]
     if len(env_policies) == 0:
         raise RuntimeError("No policies were provided!")
     for env, policy, extra_overrides, file_name_extra in env_policies:
+        print(f"env={env} policy={policy} extra_overrides={extra_overrides}")
         save_policy_to_safetensors(
             "tests/data/save_policy_to_safetensors", env, policy, extra_overrides, file_name_extra
         )
+        print()

tests/test_policies.py

@@ -147,10 +147,11 @@ def test_policy(env_name, policy_name, extra_overrides):
     # Check that we run select_actions and get the appropriate output.
     env = make_env(cfg, n_envs=2)
 
+    batch_size = 2
     dataloader = torch.utils.data.DataLoader(
         dataset,
         num_workers=0,
-        batch_size=2,
+        batch_size=batch_size,
         shuffle=True,
         pin_memory=DEVICE != "cpu",
         drop_last=True,
@@ -164,12 +165,19 @@ def test_policy(env_name, policy_name, extra_overrides):
 
     # Test updating the policy (and test that it does not mutate the batch)
     batch_ = deepcopy(batch)
-    policy.forward(batch)
+    out = policy.forward(batch)
     assert set(batch) == set(batch_), "Batch keys are not the same after a forward pass."
     assert all(
         torch.equal(batch[k], batch_[k]) for k in batch
     ), "Batch values are not the same after a forward pass."
 
+    # Test loss can be visualized using visualize_dataset_html.py
+    for key in out:
+        if "loss" in key:
+            assert (
+                out[key].ndim == 1 and out[key].shape[0] == batch_size
+            ), f"1 loss value per item in the batch is expected, but {out[key].shape} provided instead."
+
     # reset the policy and environment
     policy.reset()
     observation, _ = env.reset(seed=cfg.seed)
@@ -234,6 +242,7 @@ def test_policy_defaults(policy_name: str):
     [
         ("xarm", "tdmpc"),
         ("pusht", "diffusion"),
+        ("pusht", "vqbet"),
         ("aloha", "act"),
     ],
 )
@@ -250,7 +259,7 @@ def test_yaml_matches_dataclass(env_name: str, policy_name: str):
 def test_save_and_load_pretrained(policy_name: str):
     policy_cls, _ = get_policy_and_config_classes(policy_name)
     policy: Policy = policy_cls()
-    save_dir = "/tmp/test_save_and_load_pretrained_{policy_cls.__name__}"
+    save_dir = f"/tmp/test_save_and_load_pretrained_{policy_cls.__name__}"
     policy.save_pretrained(save_dir)
     policy_ = policy_cls.from_pretrained(save_dir)
     assert all(torch.equal(p, p_) for p, p_ in zip(policy.parameters(), policy_.parameters(), strict=True))
@@ -365,6 +374,7 @@ def test_normalize(insert_temporal_dim):
             ["policy.n_action_steps=8", "policy.num_inference_steps=10", "policy.down_dims=[128, 256, 512]"],
             "",
         ),
+        ("pusht", "vqbet", "[]", ""),
         ("aloha", "act", ["policy.n_action_steps=10"], ""),
         ("aloha", "act", ["policy.n_action_steps=1000", "policy.chunk_size=1000"], "_1000_steps"),
         ("dora_aloha_real", "act_real", ["policy.n_action_steps=10"], ""),
@@ -461,7 +471,3 @@ def test_act_temporal_ensembler():
         assert torch.all(offline_avg <= einops.reduce(seq_slice, "b s 1 -> b 1", "max"))
         # Selected atol=1e-4 keeping in mind actions in [-1, 1] and excepting 0.01% error.
         assert torch.allclose(online_avg, offline_avg, atol=1e-4)
-
-
-if __name__ == "__main__":
-    test_act_temporal_ensembler()

tests/test_visualize_dataset.py

@@ -25,13 +25,13 @@ from lerobot.scripts.visualize_dataset import visualize_dataset
     ["lerobot/pusht"],
 )
 @pytest.mark.parametrize("root", [Path(__file__).parent / "data"])
-def test_visualize_local_dataset(tmpdir, repo_id, root):
+def test_visualize_dataset_root(tmpdir, repo_id, root):
     rrd_path = visualize_dataset(
         repo_id,
+        root=root,
         episode_index=0,
         batch_size=32,
         save=True,
         output_dir=tmpdir,
-        root=root,
     )
     assert rrd_path.exists()