Remove EMA model from Diffusion Policy (#134)

tests/scripts/save_policy_to_safetensor.py

@@ -88,14 +88,8 @@ def save_policy_to_safetensors(output_dir, env_name, policy_name, extra_override
 
 
 if __name__ == "__main__":
-    env_policies = [
-        # ("xarm", "tdmpc", ["policy.n_action_repeats=2"]),
-        (
-            "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"]),
-    ]
+    # Instructions: include the policies that you want to save artifacts for here. Please make sure to revert
+    # your changes when you are done.
+    env_policies = []
     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

@@ -249,6 +249,17 @@ def test_normalize(insert_temporal_dim):
 # pass if it's run on another platform due to floating point errors
 @require_x86_64_kernel
 def test_backward_compatibility(env_name, policy_name, extra_overrides):
+    """
+    NOTE: If this test does not pass, and you have intentionally changed something in the policy:
+        1. Inspect the differences in policy outputs and make sure you can account for them. Your PR should
+           include a report on what changed and how that affected the outputs.
+        2. Go to the `if __name__ == "__main__"` block of `test/scripts/save_policy_to_safetensors.py` and
+           add the policies you want to update the test artifacts for.
+        3. Run `python test/scripts/save_policy_to_safetensors.py`. The test artifact should be updated.
+        4. Check that this test now passes.
+        5. Remember to restore `test/scripts/save_policy_to_safetensors.py` to its original state.
+        6. Remember to stage and commit the resulting changes to `tests/data`.
+    """
     env_policy_dir = Path("tests/data/save_policy_to_safetensors") / f"{env_name}_{policy_name}"
     saved_output_dict = load_file(env_policy_dir / "output_dict.safetensors")
     saved_grad_stats = load_file(env_policy_dir / "grad_stats.safetensors")