From 30a2392c73ab811ff07e2394d87bbeb12851a4f6 Mon Sep 17 00:00:00 2001 From: sean1295 Date: Wed, 23 Jul 2025 06:29:14 -0400 Subject: [PATCH 1/5] update diffusers==0.34.0 --- robomimic/algo/diffusion_policy.py | 14 ++++++++++---- setup.py | 2 +- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/robomimic/algo/diffusion_policy.py b/robomimic/algo/diffusion_policy.py index f1a87e79..e1d4d025 100644 --- a/robomimic/algo/diffusion_policy.py +++ b/robomimic/algo/diffusion_policy.py @@ -109,13 +109,16 @@ def _create_networks(self): # setup EMA ema = None + ema_nets = None if self.algo_config.ema.enabled: - ema = EMAModel(model=nets, power=self.algo_config.ema.power) + ema = EMAModel(parameters=nets.parameters(), power=self.algo_config.ema.power) + ema_nets = deepcopy(nets) # set attrs self.nets = nets self.noise_scheduler = noise_scheduler self.ema = ema + self.ema_nets = ema_nets self.action_check_done = False self.obs_queue = None self.action_queue = None @@ -232,7 +235,7 @@ def train_on_batch(self, batch, epoch, validate=False): # update Exponential Moving Average of the model weights if self.ema is not None: - self.ema.step(self.nets) + self.ema.step(self.nets.parameters()) step_info = { "policy_grad_norms": policy_grad_norms @@ -330,7 +333,7 @@ def _get_action_trajectory(self, obs_dict, goal_dict=None): # select network nets = self.nets if self.ema is not None: - nets = self.ema.averaged_model + nets = self.ema_nets # encode obs inputs = { @@ -382,6 +385,8 @@ def serialize(self): """ Get dictionary of current model parameters. """ + if self.ema is not None: + self.ema.copy_to(self.ema_nets.parameters()) return { "nets": self.nets.state_dict(), "optimizers": { k : self.optimizers[k].state_dict() for k in self.optimizers }, @@ -408,7 +413,8 @@ def deserialize(self, model_dict, load_optimizers=False): model_dict["lr_schedulers"] = {} if model_dict.get("ema", None) is not None: - self.ema.averaged_model.load_state_dict(model_dict["ema"]) + self.ema.load_state_dict(model_dict["ema"]) + self.ema_nets.load_state_dict(model_dict["nets"]) if load_optimizers: for k in model_dict["optimizers"]: diff --git a/setup.py b/setup.py index 704db21b..7bc8cc11 100644 --- a/setup.py +++ b/setup.py @@ -31,7 +31,7 @@ "torchvision", "huggingface_hub==0.23.4", "transformers==4.41.2", - "diffusers==0.11.1", + "diffusers==0.34.0", ], eager_resources=['*'], include_package_data=True, From af07b27d6b98e4f9ac2f2d2ad1771e8ecaca1dba Mon Sep 17 00:00:00 2001 From: sean1295 Date: Wed, 23 Jul 2025 06:30:36 -0400 Subject: [PATCH 2/5] update diffusers==0.34.0 --- robomimic/algo/diffusion_policy.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/robomimic/algo/diffusion_policy.py b/robomimic/algo/diffusion_policy.py index e1d4d025..8d5db092 100644 --- a/robomimic/algo/diffusion_policy.py +++ b/robomimic/algo/diffusion_policy.py @@ -2,10 +2,9 @@ Implementation of Diffusion Policy https://diffusion-policy.cs.columbia.edu/ by Cheng Chi """ from typing import Callable, Union -import math +from copy import deepcopy from collections import OrderedDict, deque from packaging.version import parse as parse_version -import random import torch import torch.nn as nn import torch.nn.functional as F From 460185d5ff17deb551125ce1afef20baa16541a3 Mon Sep 17 00:00:00 2001 From: sean1295 Date: Wed, 23 Jul 2025 06:31:04 -0400 Subject: [PATCH 3/5] remove comment --- robomimic/algo/diffusion_policy.py | 1 - 1 file changed, 1 deletion(-) diff --git a/robomimic/algo/diffusion_policy.py b/robomimic/algo/diffusion_policy.py index 8d5db092..2b18a1ef 100644 --- a/robomimic/algo/diffusion_policy.py +++ b/robomimic/algo/diffusion_policy.py @@ -8,7 +8,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -# requires diffusers==0.11.1 from diffusers.schedulers.scheduling_ddpm import DDPMScheduler from diffusers.schedulers.scheduling_ddim import DDIMScheduler from diffusers.training_utils import EMAModel From d13b21fe5ab6b15c688dad111af1610782cd330b Mon Sep 17 00:00:00 2001 From: sean1295 Date: Thu, 31 Jul 2025 00:48:03 -0400 Subject: [PATCH 4/5] remove version requirement for huggingface_hub --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 7bc8cc11..0d16fcb3 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ "egl_probe>=1.0.1", "torch", "torchvision", - "huggingface_hub==0.23.4", + "huggingface_hub", "transformers==4.41.2", "diffusers==0.34.0", ], From a2837deeb72a54de2401f4f00d1b448c5255b7ca Mon Sep 17 00:00:00 2001 From: sean1295 Date: Fri, 1 Aug 2025 11:56:55 -0400 Subject: [PATCH 5/5] always force eval mode for ema model --- robomimic/algo/diffusion_policy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/robomimic/algo/diffusion_policy.py b/robomimic/algo/diffusion_policy.py index 2b18a1ef..84077cc2 100644 --- a/robomimic/algo/diffusion_policy.py +++ b/robomimic/algo/diffusion_policy.py @@ -20,7 +20,6 @@ from robomimic.algo import register_algo_factory_func, PolicyAlgo -import random import robomimic.utils.torch_utils as TorchUtils import robomimic.utils.tensor_utils as TensorUtils import robomimic.utils.obs_utils as ObsUtils @@ -117,6 +116,7 @@ def _create_networks(self): self.noise_scheduler = noise_scheduler self.ema = ema self.ema_nets = ema_nets + self.ema_nets.eval() self.action_check_done = False self.obs_queue = None self.action_queue = None