More frequent EMA update (#20)

* move ema update within pretraining epoch * update pretraining ema configs * add lift and can epoch 8000 checkpoint url * add note about EMA issue in pretraining instruction
2024-11-06 20:42:31 -05:00 · 2024-11-06 20:42:31 -05:00 · e1ef4ca1cf
commit e1ef4ca1cf
parent dc8e0c9edc
91 changed files with 183 additions and 171 deletions
--- a/agent/pretrain/train_diffusion_agent.py
+++ b/agent/pretrain/train_diffusion_agent.py
@ -36,6 +36,10 @@ class TrainDiffusionAgent(PreTrainAgent):

                self.optimizer.step()
                self.optimizer.zero_grad()
+
+                # update ema
+                if self.epoch % self.update_ema_freq == 0:
+                    self.step_ema()
            loss_train = np.mean(loss_train_epoch)

            # validate
@ -53,10 +57,6 @@ class TrainDiffusionAgent(PreTrainAgent):
            # update lr
            self.lr_scheduler.step()

-            # update ema
-            if self.epoch % self.update_ema_freq == 0:
-                self.step_ema()
-
            # save model
            if self.epoch % self.save_model_freq == 0 or self.epoch == self.n_epochs:
                self.save_model()
--- a/agent/pretrain/train_gaussian_agent.py
+++ b/agent/pretrain/train_gaussian_agent.py
@ -44,6 +44,10 @@ class TrainGaussianAgent(PreTrainAgent):

                self.optimizer.step()
                self.optimizer.zero_grad()
+
+                # update ema
+                if self.epoch % self.update_ema_freq == 0:
+                    self.step_ema()
            loss_train = np.mean(loss_train_epoch)
            ent_train = np.mean(ent_train_epoch)

@ -65,10 +69,6 @@ class TrainGaussianAgent(PreTrainAgent):
            # update lr
            self.lr_scheduler.step()

-            # update ema
-            if self.epoch % self.update_ema_freq == 0:
-                self.step_ema()
-
            # save model
            if self.epoch % self.save_model_freq == 0 or self.epoch == self.n_epochs:
                self.save_model()
--- a/cfg/d3il/pretrain/avoid_m1/pre_diffusion_mlp.yaml
+++ b/cfg/d3il/pretrain/avoid_m1/pre_diffusion_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/d3il/pretrain/avoid_m1/pre_gaussian_mlp.yaml
+++ b/cfg/d3il/pretrain/avoid_m1/pre_gaussian_mlp.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/d3il/pretrain/avoid_m1/pre_gmm_mlp.yaml
+++ b/cfg/d3il/pretrain/avoid_m1/pre_gmm_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gmm.GMMModel
--- a/cfg/d3il/pretrain/avoid_m2/pre_diffusion_mlp.yaml
+++ b/cfg/d3il/pretrain/avoid_m2/pre_diffusion_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/d3il/pretrain/avoid_m2/pre_gaussian_mlp.yaml
+++ b/cfg/d3il/pretrain/avoid_m2/pre_gaussian_mlp.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/d3il/pretrain/avoid_m2/pre_gmm_mlp.yaml
+++ b/cfg/d3il/pretrain/avoid_m2/pre_gmm_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gmm.GMMModel
--- a/cfg/d3il/pretrain/avoid_m3/pre_diffusion_mlp.yaml
+++ b/cfg/d3il/pretrain/avoid_m3/pre_diffusion_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/d3il/pretrain/avoid_m3/pre_gaussian_mlp.yaml
+++ b/cfg/d3il/pretrain/avoid_m3/pre_gaussian_mlp.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/d3il/pretrain/avoid_m3/pre_gmm_mlp.yaml
+++ b/cfg/d3il/pretrain/avoid_m3/pre_gmm_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gmm.GMMModel
--- a/cfg/furniture/pretrain/lamp_low/pre_diffusion_mlp.yaml
+++ b/cfg/furniture/pretrain/lamp_low/pre_diffusion_mlp.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/lamp_low/pre_diffusion_unet.yaml
+++ b/cfg/furniture/pretrain/lamp_low/pre_diffusion_unet.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/lamp_low/pre_gaussian_mlp.yaml
+++ b/cfg/furniture/pretrain/lamp_low/pre_gaussian_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/furniture/pretrain/lamp_med/pre_diffusion_mlp.yaml
+++ b/cfg/furniture/pretrain/lamp_med/pre_diffusion_mlp.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/lamp_med/pre_diffusion_unet.yaml
+++ b/cfg/furniture/pretrain/lamp_med/pre_diffusion_unet.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/lamp_med/pre_gaussian_mlp.yaml
+++ b/cfg/furniture/pretrain/lamp_med/pre_gaussian_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/furniture/pretrain/one_leg_low/pre_diffusion_mlp.yaml
+++ b/cfg/furniture/pretrain/one_leg_low/pre_diffusion_mlp.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/one_leg_low/pre_diffusion_unet.yaml
+++ b/cfg/furniture/pretrain/one_leg_low/pre_diffusion_unet.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/one_leg_low/pre_gaussian_mlp.yaml
+++ b/cfg/furniture/pretrain/one_leg_low/pre_gaussian_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/furniture/pretrain/one_leg_med/pre_diffusion_mlp.yaml
+++ b/cfg/furniture/pretrain/one_leg_med/pre_diffusion_mlp.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/one_leg_med/pre_diffusion_unet.yaml
+++ b/cfg/furniture/pretrain/one_leg_med/pre_diffusion_unet.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/one_leg_med/pre_gaussian_mlp.yaml
+++ b/cfg/furniture/pretrain/one_leg_med/pre_gaussian_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/furniture/pretrain/round_table_low/pre_diffusion_mlp.yaml
+++ b/cfg/furniture/pretrain/round_table_low/pre_diffusion_mlp.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/round_table_low/pre_diffusion_unet.yaml
+++ b/cfg/furniture/pretrain/round_table_low/pre_diffusion_unet.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/round_table_low/pre_gaussian_mlp.yaml
+++ b/cfg/furniture/pretrain/round_table_low/pre_gaussian_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/furniture/pretrain/round_table_med/pre_diffusion_mlp.yaml
+++ b/cfg/furniture/pretrain/round_table_med/pre_diffusion_mlp.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/round_table_med/pre_diffusion_unet.yaml
+++ b/cfg/furniture/pretrain/round_table_med/pre_diffusion_unet.yaml
@ -35,8 +35,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/furniture/pretrain/round_table_med/pre_gaussian_mlp.yaml
+++ b/cfg/furniture/pretrain/round_table_med/pre_gaussian_mlp.yaml
@ -34,8 +34,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/gym/pretrain/halfcheetah-medium-v2/pre_diffusion_mlp.yaml
+++ b/cfg/gym/pretrain/halfcheetah-medium-v2/pre_diffusion_mlp.yaml
@ -32,8 +32,8 @@ train:
    first_cycle_steps: 3000
    warmup_steps: 1
    min_lr: 1e-4
-  epoch_start_ema: 10
-  update_ema_freq: 5
+  epoch_start_ema: 20
+  update_ema_freq: 1
  save_model_freq: 100

 model:
--- a/cfg/gym/pretrain/halfcheetah-medium-v2/pre_gaussian_mlp.yaml
+++ b/cfg/gym/pretrain/halfcheetah-medium-v2/pre_gaussian_mlp.yaml
@ -31,8 +31,8 @@ train:
    first_cycle_steps: 1000
    warmup_steps: 1
    min_lr: 1e-4
-  epoch_start_ema: 10
-  update_ema_freq: 5
+  epoch_start_ema: 20
+  update_ema_freq: 1
  save_model_freq: 100

 model:
--- a/cfg/gym/pretrain/hopper-medium-v2/pre_diffusion_mlp.yaml
+++ b/cfg/gym/pretrain/hopper-medium-v2/pre_diffusion_mlp.yaml
@ -32,8 +32,8 @@ train:
    first_cycle_steps: 3000
    warmup_steps: 1
    min_lr: 1e-4
-  epoch_start_ema: 10
-  update_ema_freq: 5
+  epoch_start_ema: 20
+  update_ema_freq: 1
  save_model_freq: 100

 model:
--- a/cfg/gym/pretrain/hopper-medium-v2/pre_gaussian_mlp.yaml
+++ b/cfg/gym/pretrain/hopper-medium-v2/pre_gaussian_mlp.yaml
@ -31,8 +31,8 @@ train:
    first_cycle_steps: 1000
    warmup_steps: 1
    min_lr: 1e-4
-  epoch_start_ema: 10
-  update_ema_freq: 5
+  epoch_start_ema: 20
+  update_ema_freq: 1
  save_model_freq: 100

 model:
--- a/cfg/gym/pretrain/kitchen-complete-v0/pre_diffusion_mlp.yaml
+++ b/cfg/gym/pretrain/kitchen-complete-v0/pre_diffusion_mlp.yaml
@ -32,9 +32,9 @@ train:
    first_cycle_steps: 8000
    warmup_steps: 1
    min_lr: 1e-4
-  epoch_start_ema: 10
-  update_ema_freq: 5
-  save_model_freq: 1000
+  epoch_start_ema: 20
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/gym/pretrain/kitchen-complete-v0/pre_gaussian_mlp.yaml
+++ b/cfg/gym/pretrain/kitchen-complete-v0/pre_gaussian_mlp.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-4
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/gym/pretrain/kitchen-mixed-v0/pre_diffusion_mlp.yaml
+++ b/cfg/gym/pretrain/kitchen-mixed-v0/pre_diffusion_mlp.yaml
@ -32,9 +32,9 @@ train:
    first_cycle_steps: 8000
    warmup_steps: 1
    min_lr: 1e-4
-  epoch_start_ema: 10
-  update_ema_freq: 5
-  save_model_freq: 1000
+  epoch_start_ema: 20
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/gym/pretrain/kitchen-mixed-v0/pre_gaussian_mlp.yaml
+++ b/cfg/gym/pretrain/kitchen-mixed-v0/pre_gaussian_mlp.yaml
@ -31,9 +31,9 @@ train:
    first_cycle_steps: 5000
    warmup_steps: 1
    min_lr: 1e-4
-  epoch_start_ema: 10
-  update_ema_freq: 5
-  save_model_freq: 1000
+  epoch_start_ema: 20
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/gym/pretrain/kitchen-partial-v0/pre_diffusion_mlp.yaml
+++ b/cfg/gym/pretrain/kitchen-partial-v0/pre_diffusion_mlp.yaml
@ -32,9 +32,9 @@ train:
    first_cycle_steps: 8000
    warmup_steps: 1
    min_lr: 1e-4
-  epoch_start_ema: 10
-  update_ema_freq: 5
-  save_model_freq: 1000
+  epoch_start_ema: 20
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/gym/pretrain/kitchen-partial-v0/pre_gaussian_mlp.yaml
+++ b/cfg/gym/pretrain/kitchen-partial-v0/pre_gaussian_mlp.yaml
@ -31,9 +31,9 @@ train:
    first_cycle_steps: 5000
    warmup_steps: 1
    min_lr: 1e-4
-  epoch_start_ema: 10
-  update_ema_freq: 5
-  save_model_freq: 1000
+  epoch_start_ema: 20
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/gym/pretrain/walker2d-medium-v2/pre_diffusion_mlp.yaml
+++ b/cfg/gym/pretrain/walker2d-medium-v2/pre_diffusion_mlp.yaml
@ -32,8 +32,8 @@ train:
    first_cycle_steps: 3000
    warmup_steps: 1
    min_lr: 1e-4
-  epoch_start_ema: 10
-  update_ema_freq: 5
+  epoch_start_ema: 20
+  update_ema_freq: 1
  save_model_freq: 100

 model:
--- a/cfg/gym/pretrain/walker2d-medium-v2/pre_gaussian_mlp.yaml
+++ b/cfg/gym/pretrain/walker2d-medium-v2/pre_gaussian_mlp.yaml
@ -31,8 +31,8 @@ train:
    first_cycle_steps: 3000
    warmup_steps: 1
    min_lr: 1e-4
-  epoch_start_ema: 10
-  update_ema_freq: 5
+  epoch_start_ema: 20
+  update_ema_freq: 1
  save_model_freq: 100

 model:
--- a/cfg/pretraining.md
+++ b/cfg/pretraining.md
@ -1,5 +1,7 @@
 ## Pre-training experiments

+**Update, Nov 6 2024**: we fixed the issue of EMA update being too infrequent causing slow pre-training. Now the number of epochs needed for pre-training can be much slower than those used in the configs. We recommend training with fewer epochs and testing the early checkpoints.
+
 ### Comparing diffusion-based RL algorithms (Sec. 5.1)
 Gym configs are under `cfg/gym/pretrain/<env_name>/`, and the config name is `pre_diffusion_mlp`. Robomimic configs are under `cfg/robomimic/pretrain/<env_name>/`, and the name is also `pre_diffusion_mlp`.

--- a/cfg/robomimic/finetune/can/ibrl_mlp.yaml
+++ b/cfg/robomimic/finetune/can/ibrl_mlp.yaml
@ -62,7 +62,7 @@ train:
    first_cycle_steps: 1000
    warmup_steps: 10
    min_lr: 1e-4
-  save_model_freq: 100000
+  save_model_freq: 50000
  val_freq: 10000
  render:
    freq: 10000
--- a/cfg/robomimic/finetune/can/ibrl_mlp_ph.yaml
+++ b/cfg/robomimic/finetune/can/ibrl_mlp_ph.yaml
@ -62,7 +62,7 @@ train:
    first_cycle_steps: 1000
    warmup_steps: 10
    min_lr: 1e-4
-  save_model_freq: 100000
+  save_model_freq: 50000
  val_freq: 10000
  render:
    freq: 10000
--- a/cfg/robomimic/finetune/square/ibrl_mlp.yaml
+++ b/cfg/robomimic/finetune/square/ibrl_mlp.yaml
@ -62,7 +62,7 @@ train:
    first_cycle_steps: 1000
    warmup_steps: 10
    min_lr: 1e-4
-  save_model_freq: 100000
+  save_model_freq: 50000
  val_freq: 10000
  render:
    freq: 10000
--- a/cfg/robomimic/finetune/square/ibrl_mlp_ph.yaml
+++ b/cfg/robomimic/finetune/square/ibrl_mlp_ph.yaml
@ -62,7 +62,7 @@ train:
    first_cycle_steps: 1000
    warmup_steps: 10
    min_lr: 1e-4
-  save_model_freq: 100000
+  save_model_freq: 50000
  val_freq: 10000
  render:
    freq: 10000
--- a/cfg/robomimic/pretrain/can/pre_diffusion_mlp.yaml
+++ b/cfg/robomimic/pretrain/can/pre_diffusion_mlp.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/can/pre_diffusion_mlp_img.yaml
+++ b/cfg/robomimic/pretrain/can/pre_diffusion_mlp_img.yaml
@ -43,7 +43,7 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
+  update_ema_freq: 1
  save_model_freq: 500

 model:
--- a/cfg/robomimic/pretrain/can/pre_diffusion_mlp_ta1.yaml
+++ b/cfg/robomimic/pretrain/can/pre_diffusion_mlp_ta1.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/can/pre_diffusion_mlp_ta1_ph.yaml
+++ b/cfg/robomimic/pretrain/can/pre_diffusion_mlp_ta1_ph.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/can/pre_diffusion_unet.yaml
+++ b/cfg/robomimic/pretrain/can/pre_diffusion_unet.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/can/pre_gaussian_mlp.yaml
+++ b/cfg/robomimic/pretrain/can/pre_gaussian_mlp.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/can/pre_gaussian_mlp_ibrl.yaml
+++ b/cfg/robomimic/pretrain/can/pre_gaussian_mlp_ibrl.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-4
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/can/pre_gaussian_mlp_img.yaml
+++ b/cfg/robomimic/pretrain/can/pre_gaussian_mlp_img.yaml
@ -42,7 +42,7 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
+  update_ema_freq: 1
  save_model_freq: 500

 model:
--- a/cfg/robomimic/pretrain/can/pre_gaussian_mlp_ta1_ph.yaml
+++ b/cfg/robomimic/pretrain/can/pre_gaussian_mlp_ta1_ph.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/can/pre_gaussian_transformer.yaml
+++ b/cfg/robomimic/pretrain/can/pre_gaussian_transformer.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/can/pre_gmm_mlp.yaml
+++ b/cfg/robomimic/pretrain/can/pre_gmm_mlp.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gmm.GMMModel
--- a/cfg/robomimic/pretrain/can/pre_gmm_transformer.yaml
+++ b/cfg/robomimic/pretrain/can/pre_gmm_transformer.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gmm.GMMModel
--- a/cfg/robomimic/pretrain/lift/pre_diffusion_mlp.yaml
+++ b/cfg/robomimic/pretrain/lift/pre_diffusion_mlp.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/lift/pre_diffusion_mlp_img.yaml
+++ b/cfg/robomimic/pretrain/lift/pre_diffusion_mlp_img.yaml
@ -43,7 +43,7 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
+  update_ema_freq: 1
  save_model_freq: 500

 model:
--- a/cfg/robomimic/pretrain/lift/pre_diffusion_unet.yaml
+++ b/cfg/robomimic/pretrain/lift/pre_diffusion_unet.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/lift/pre_gaussian_mlp.yaml
+++ b/cfg/robomimic/pretrain/lift/pre_gaussian_mlp.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/lift/pre_gaussian_mlp_img.yaml
+++ b/cfg/robomimic/pretrain/lift/pre_gaussian_mlp_img.yaml
@ -42,7 +42,7 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
+  update_ema_freq: 1
  save_model_freq: 500

 model:
--- a/cfg/robomimic/pretrain/lift/pre_gaussian_transformer.yaml
+++ b/cfg/robomimic/pretrain/lift/pre_gaussian_transformer.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/lift/pre_gmm_mlp.yaml
+++ b/cfg/robomimic/pretrain/lift/pre_gmm_mlp.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gmm.GMMModel
--- a/cfg/robomimic/pretrain/lift/pre_gmm_transformer.yaml
+++ b/cfg/robomimic/pretrain/lift/pre_gmm_transformer.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gmm.GMMModel
--- a/cfg/robomimic/pretrain/square/pre_diffusion_mlp.yaml
+++ b/cfg/robomimic/pretrain/square/pre_diffusion_mlp.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/square/pre_diffusion_mlp_img.yaml
+++ b/cfg/robomimic/pretrain/square/pre_diffusion_mlp_img.yaml
@ -43,7 +43,7 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
+  update_ema_freq: 1
  save_model_freq: 500

 model:
--- a/cfg/robomimic/pretrain/square/pre_diffusion_mlp_ta1.yaml
+++ b/cfg/robomimic/pretrain/square/pre_diffusion_mlp_ta1.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/square/pre_diffusion_mlp_ta1_ph.yaml
+++ b/cfg/robomimic/pretrain/square/pre_diffusion_mlp_ta1_ph.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/square/pre_diffusion_unet.yaml
+++ b/cfg/robomimic/pretrain/square/pre_diffusion_unet.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/square/pre_gaussian_mlp.yaml
+++ b/cfg/robomimic/pretrain/square/pre_gaussian_mlp.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/square/pre_gaussian_mlp_ibrl.yaml
+++ b/cfg/robomimic/pretrain/square/pre_gaussian_mlp_ibrl.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-4
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/square/pre_gaussian_mlp_img.yaml
+++ b/cfg/robomimic/pretrain/square/pre_gaussian_mlp_img.yaml
@ -42,7 +42,7 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
+  update_ema_freq: 1
  save_model_freq: 500

 model:
--- a/cfg/robomimic/pretrain/square/pre_gaussian_mlp_ta1_ph.yaml
+++ b/cfg/robomimic/pretrain/square/pre_gaussian_mlp_ta1_ph.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/square/pre_gaussian_transformer.yaml
+++ b/cfg/robomimic/pretrain/square/pre_gaussian_transformer.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/square/pre_gmm_mlp.yaml
+++ b/cfg/robomimic/pretrain/square/pre_gmm_mlp.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gmm.GMMModel
--- a/cfg/robomimic/pretrain/square/pre_gmm_transformer.yaml
+++ b/cfg/robomimic/pretrain/square/pre_gmm_transformer.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gmm.GMMModel
--- a/cfg/robomimic/pretrain/transport/pre_diffusion_mlp.yaml
+++ b/cfg/robomimic/pretrain/transport/pre_diffusion_mlp.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/transport/pre_diffusion_mlp_img.yaml
+++ b/cfg/robomimic/pretrain/transport/pre_diffusion_mlp_img.yaml
@ -43,7 +43,7 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
+  update_ema_freq: 1
  save_model_freq: 500

 model:
--- a/cfg/robomimic/pretrain/transport/pre_diffusion_unet.yaml
+++ b/cfg/robomimic/pretrain/transport/pre_diffusion_unet.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.diffusion.diffusion.DiffusionModel
--- a/cfg/robomimic/pretrain/transport/pre_gaussian_mlp.yaml
+++ b/cfg/robomimic/pretrain/transport/pre_gaussian_mlp.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/transport/pre_gaussian_mlp_img.yaml
+++ b/cfg/robomimic/pretrain/transport/pre_gaussian_mlp_img.yaml
@ -42,7 +42,7 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
+  update_ema_freq: 1
  save_model_freq: 500

 model:
--- a/cfg/robomimic/pretrain/transport/pre_gaussian_transformer.yaml
+++ b/cfg/robomimic/pretrain/transport/pre_gaussian_transformer.yaml
@ -32,8 +32,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gaussian.GaussianModel
--- a/cfg/robomimic/pretrain/transport/pre_gmm_mlp.yaml
+++ b/cfg/robomimic/pretrain/transport/pre_gmm_mlp.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gmm.GMMModel
--- a/cfg/robomimic/pretrain/transport/pre_gmm_transformer.yaml
+++ b/cfg/robomimic/pretrain/transport/pre_gmm_transformer.yaml
@ -33,8 +33,8 @@ train:
    warmup_steps: 100
    min_lr: 1e-5
  epoch_start_ema: 20
-  update_ema_freq: 10
-  save_model_freq: 1000
+  update_ema_freq: 1
+  save_model_freq: 500

 model:
  _target_: model.common.gmm.GMMModel
--- a/cfg/robomimic/scratch/can/rlpd_mlp.yaml
+++ b/cfg/robomimic/scratch/can/rlpd_mlp.yaml
@ -61,7 +61,7 @@ train:
    first_cycle_steps: 1000
    warmup_steps: 10
    min_lr: 1e-4
-  save_model_freq: 100000
+  save_model_freq: 50000
  val_freq: 10000
  render:
    freq: 10000
--- a/cfg/robomimic/scratch/can/rlpd_mlp_ph.yaml
+++ b/cfg/robomimic/scratch/can/rlpd_mlp_ph.yaml
@ -61,7 +61,7 @@ train:
    first_cycle_steps: 1000
    warmup_steps: 10
    min_lr: 1e-4
-  save_model_freq: 100000
+  save_model_freq: 50000
  val_freq: 10000
  render:
    freq: 10000
--- a/cfg/robomimic/scratch/square/rlpd_mlp.yaml
+++ b/cfg/robomimic/scratch/square/rlpd_mlp.yaml
@ -61,7 +61,7 @@ train:
    first_cycle_steps: 1000
    warmup_steps: 10
    min_lr: 1e-4
-  save_model_freq: 100000
+  save_model_freq: 50000
  val_freq: 10000
  render:
    freq: 10000
--- a/cfg/robomimic/scratch/square/rlpd_mlp_ph.yaml
+++ b/cfg/robomimic/scratch/square/rlpd_mlp_ph.yaml
@ -61,7 +61,7 @@ train:
    first_cycle_steps: 1000
    warmup_steps: 10
    min_lr: 1e-4
-  save_model_freq: 100000
+  save_model_freq: 50000
  val_freq: 10000
  render:
    freq: 10000
--- a/script/download_url.py
+++ b/script/download_url.py
@ -284,6 +284,11 @@ def get_checkpoint_download_url(cfg):
        in path
    ):
        return "https://drive.google.com/file/d/1Ngr-DNxoB9XNCZ2O-NF5p60NzmYlzmWG/view?usp=drive_link"
+    elif (
+        "lift_pre_diffusion_mlp_ta4_td20/2024-06-28_14-47-58/checkpoint/state_8000.pt"
+        in path
+    ):
+        return "https://drive.google.com/file/d/1IyXa6CEXO16mmCCHgNfFTnmvAhA3PVxQ/view?usp=drive_link"
    elif (
        "lift_pre_diffusion_mlp_img_ta4_td100/2024-07-30_22-24-35/checkpoint/state_2500.pt"
        in path
@ -323,6 +328,11 @@ def get_checkpoint_download_url(cfg):
        in path
    ):
        return "https://drive.google.com/file/d/1L1ZLD1u1Y1YJmRLGzScXbQ02wGS-_cWo/view?usp=drive_link"
+    elif (
+        "can_pre_diffusion_mlp_ta4_td20/2024-06-28_13-29-54/checkpoint/state_8000.pt"
+        in path
+    ):
+        return "https://drive.google.com/file/d/1_3-QcDrWCH6cPRPLuVnYQt25ymvBYHgn/view?usp=drive_link"
    elif (
        "can_pre_diffusion_mlp_img_ta4_td100/2024-07-30_22-23-55/checkpoint/state_5000.pt"
        in path