Remove EMA model from Diffusion Policy #134

alexander-soare · 2024-05-05T07:51:03Z

What does this PR do?

As the title suggests. Also updates test artifacts for diffusion policy backwards compatibility check.

Side change:

Adds documentation for workflow required to update the test artifacts for the policy b/c checks.

How was it tested?

I have a pretrained diffusion policy that reaches SOTA eval metrics. I evaluated it for 500 episodes using the EMA vs non-EMA weights.

{
  "ema": {
    "avg_sum_reward": 104.15279900893682,
    "avg_max_reward": 0.9495501381835949,
    "pc_success": 64.4,
    "eval_s": 570.6887876987457,
    "eval_ep_s": 1.1413775758743285
  }
  "non_ema": {
    "avg_sum_reward": 108.0149595484424,
    "avg_max_reward": 0.9588220574225956,
    "pc_success": 63.800000000000004,
    "eval_s": 555.2821447849274,
    "eval_ep_s": 1.1105642900466919
  },
}

The mean "avg_max_reward" is higher for non-EMA (without considering error-bars). For success rate we can take a uniform prior and calculate the posterior beta distribution to get the mean, upper confidence bound (mean + 34.1%) and lower confidence bound (mean - 34.1%)

import scipy.stats

num_episodes = 500
success_rate = 0.644  # or 0.638 for non-ema

alpha = num_episodes * success_rate + 1
beta = num_episodes * (1 - success_rate) + 1

confidence_interval = 0.682
lower_percentile = (1 - confidence_interval) / 2
upper_percentile = 1 - lower_percentile

print("Mean:", scipy.stats.beta.mean(alpha, beta))
print("Lower:", scipy.stats.beta.ppf(lower_percentile, alpha, beta))
print("Upper:", scipy.stats.beta.ppf(upper
[ema.json](https://github.com/huggingface/lerobot/files/15212103/ema.json)
[no_ema.json](https://github.com/huggingface/lerobot/files/15212104/no_ema.json)
_percentile, alpha, beta))

EMA results:
Mean: 0.6434262948207171
Lower: 0.6220813179507085
Upper: 0.6647718949623826

Non-EMA results:
Mean: 0.6374501992031872
Lower: 0.6160270432185413
Upper: 0.6588739530055447

The means are not significantly different.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR. Try to avoid tagging more than 3 people.

Cadene

Genius! so happy we removed EMA!

Cadene · 2024-05-05T10:00:14Z

tests/scripts/save_policy_to_safetensor.py

+        # (
+        #     "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"]),


Should we remove?

I did, and updated instructions.

Cadene · 2024-05-05T10:00:24Z

tests/test_policies.py

            ["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=10"]),


Should we remove?

This should be uncommented actually. I reverted.

This should be uncommented actually. I reverted.

Was about to say the same, nice!

Cadene · 2024-05-05T10:02:38Z

lerobot/common/policies/diffusion/modeling_diffusion.py

-
-        Note: this method uses the ema model weights if self.training == False, otherwise the non-ema model
-        weights.


Could you add a note about EMA, saying that we tested with and without, and got as good or better results without EMA, so we decided to remove it for sake of simplicity?

Okay but I added them in the yaml config as this detail is more relevant to the outer scope. Ptal

aliberts

LGTM

aliberts · 2024-05-05T10:15:50Z

tests/test_policies.py

            ["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=10"]),


This should be uncommented actually. I reverted.

Was about to say the same, nice!

aliberts · 2024-05-05T10:17:05Z

tests/test_policies.py

+    """
+    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
+           comment in 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`.
+    """


I should have done that, that's really helpful, thanks!

alexander-soare added 2 commits May 5, 2024 08:18

remove EMA from DP

ee55d28

remove EMA from diffusion policy

3ecf6b4

alexander-soare added the policies Items related to robot policies label May 5, 2024

alexander-soare requested review from Cadene and aliberts May 5, 2024 07:51

Cadene approved these changes May 5, 2024

View reviewed changes

revision

2741b5c

aliberts approved these changes May 5, 2024

View reviewed changes

alexander-soare merged commit f3bba02 into huggingface:main May 5, 2024

alexander-soare deleted the remove_ema_from_dp branch May 5, 2024 10:26

aliberts mentioned this pull request May 8, 2024

Enable logging all the information returned by the forward methods of policies #151

Merged

Kimho666 mentioned this pull request Nov 8, 2024

Low accuracy for diffusion policy+aloha env+sim_transfer_cude_human dataset #502

Open

menhguin pushed a commit to menhguin/lerobot that referenced this pull request Feb 9, 2025

Remove EMA model from Diffusion Policy (huggingface#134)

0bd4f7b

Kalcy-U referenced this pull request in Kalcy-U/lerobot May 13, 2025

Remove EMA model from Diffusion Policy (#134)

cc7d2cf

ZoreAnuj pushed a commit to luckyrobots/lerobot that referenced this pull request Jul 29, 2025

Remove EMA model from Diffusion Policy (huggingface#134)

4e8951c

Ricci084 pushed a commit to JeffWang987/lerobot that referenced this pull request Sep 5, 2025

Remove EMA model from Diffusion Policy (huggingface#134)

33a5ff5

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Remove EMA model from Diffusion Policy #134

Remove EMA model from Diffusion Policy #134

Uh oh!

alexander-soare commented May 5, 2024 •

edited

Loading

Uh oh!

Cadene left a comment

Uh oh!

Cadene May 5, 2024

Uh oh!

alexander-soare May 5, 2024

Uh oh!

Cadene May 5, 2024

Uh oh!

alexander-soare May 5, 2024

Uh oh!

aliberts May 5, 2024

Uh oh!

Cadene May 5, 2024

Uh oh!

alexander-soare May 5, 2024

Uh oh!

aliberts left a comment

Uh oh!

aliberts May 5, 2024

Uh oh!

aliberts May 5, 2024

Uh oh!

Uh oh!


		Note: this method uses the ema model weights if self.training == False, otherwise the non-ema model
		weights.

Remove EMA model from Diffusion Policy #134

Remove EMA model from Diffusion Policy #134

Uh oh!

Conversation

alexander-soare commented May 5, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

How was it tested?

Who can review?

Uh oh!

Cadene left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aliberts left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

alexander-soare commented May 5, 2024 •

edited

Loading