use default epoch_start_ema=20 and update_ema_freq=10

agent/pretrain/train_agent.py

@@ -83,7 +83,8 @@ class PreTrainAgent:
         # Training params
         self.n_epochs = cfg.train.n_epochs
         self.batch_size = cfg.train.batch_size
-        self.epoch_start_ema = cfg.train.epoch_start_ema
+        self.epoch_start_ema = cfg.train.get("epoch_start_ema", 20)
+        self.update_ema_freq = cfg.train.get("update_ema_freq", 10)
         self.val_freq = cfg.train.get("val_freq", 100)
 
         # Logging, checkpoints

agent/pretrain/train_diffusion_agent.py

@@ -21,6 +21,7 @@ class TrainDiffusionAgent(PreTrainAgent):
 
         timer = Timer()
         self.epoch = 1
+        cnt_batch = 0
         for _ in range(self.n_epochs):
 
             # train
@@ -38,7 +39,9 @@ class TrainDiffusionAgent(PreTrainAgent):
                 self.optimizer.zero_grad()
 
                 # update ema
+                if cnt_batch % self.update_ema_freq == 0:
                     self.step_ema()
+                cnt_batch += 1
             loss_train = np.mean(loss_train_epoch)
 
             # validate

agent/pretrain/train_gaussian_agent.py

@@ -24,6 +24,7 @@ class TrainGaussianAgent(PreTrainAgent):
 
         timer = Timer()
         self.epoch = 1
+        cnt_batch = 0
         for _ in range(self.n_epochs):
 
             # train
@@ -46,7 +47,9 @@ class TrainGaussianAgent(PreTrainAgent):
                 self.optimizer.zero_grad()
 
                 # update ema
+                if cnt_batch % self.update_ema_freq == 0:
                     self.step_ema()
+                cnt_batch += 1
             loss_train = np.mean(loss_train_epoch)
             ent_train = np.mean(ent_train_epoch)
 

cfg/d3il/pretrain/avoid_m1/pre_diffusion_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 15000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/d3il/pretrain/avoid_m1/pre_gaussian_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/d3il/pretrain/avoid_m1/pre_gmm_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/d3il/pretrain/avoid_m2/pre_diffusion_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 15000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/d3il/pretrain/avoid_m2/pre_gaussian_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/d3il/pretrain/avoid_m2/pre_gmm_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/d3il/pretrain/avoid_m3/pre_diffusion_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 15000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/d3il/pretrain/avoid_m3/pre_gaussian_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/d3il/pretrain/avoid_m3/pre_gmm_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/lamp_low/pre_diffusion_mlp.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/lamp_low/pre_diffusion_unet.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/lamp_low/pre_gaussian_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/lamp_med/pre_diffusion_mlp.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/lamp_med/pre_diffusion_unet.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/lamp_med/pre_gaussian_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/one_leg_low/pre_diffusion_mlp.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/one_leg_low/pre_diffusion_unet.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/one_leg_low/pre_gaussian_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/one_leg_med/pre_diffusion_mlp.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/one_leg_med/pre_diffusion_unet.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/one_leg_med/pre_gaussian_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/round_table_low/pre_diffusion_mlp.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/round_table_low/pre_diffusion_unet.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/round_table_low/pre_gaussian_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/round_table_med/pre_diffusion_mlp.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/round_table_med/pre_diffusion_unet.yaml

@@ -34,7 +34,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/furniture/pretrain/round_table_med/pre_gaussian_mlp.yaml

@@ -33,7 +33,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/gym/pretrain/halfcheetah-medium-v2/pre_diffusion_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 3000
     warmup_steps: 1
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 100
 
 model:

cfg/gym/pretrain/halfcheetah-medium-v2/pre_gaussian_mlp.yaml

@@ -31,7 +31,7 @@ train:
     first_cycle_steps: 1000
     warmup_steps: 1
     min_lr: 1e-4
-  epoch_start_ema: 20
+ 
   save_model_freq: 100
 
 model:

cfg/gym/pretrain/hopper-medium-v2/pre_diffusion_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 3000
     warmup_steps: 1
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 100
 
 model:

cfg/gym/pretrain/hopper-medium-v2/pre_gaussian_mlp.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 1000
     warmup_steps: 1
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 100
 
 model:

cfg/gym/pretrain/kitchen-complete-v0/pre_diffusion_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 8000
     warmup_steps: 1
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/gym/pretrain/kitchen-complete-v0/pre_gaussian_mlp.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/gym/pretrain/kitchen-mixed-v0/pre_diffusion_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 8000
     warmup_steps: 1
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/gym/pretrain/kitchen-mixed-v0/pre_gaussian_mlp.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 1
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/gym/pretrain/kitchen-partial-v0/pre_diffusion_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 8000
     warmup_steps: 1
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/gym/pretrain/kitchen-partial-v0/pre_gaussian_mlp.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 1
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/gym/pretrain/walker2d-medium-v2/pre_diffusion_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 3000
     warmup_steps: 1
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 100
 
 model:

cfg/gym/pretrain/walker2d-medium-v2/pre_gaussian_mlp.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 3000
     warmup_steps: 1
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 100
 
 model:

cfg/robomimic/pretrain/can/pre_diffusion_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/can/pre_diffusion_mlp_img.yaml

@@ -42,7 +42,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/can/pre_diffusion_mlp_ta1.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/can/pre_diffusion_mlp_ta1_ph.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/can/pre_diffusion_unet.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/can/pre_gaussian_mlp.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/can/pre_gaussian_mlp_ibrl.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/can/pre_gaussian_mlp_img.yaml

@@ -41,7 +41,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/can/pre_gaussian_mlp_ta1_ph.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/can/pre_gaussian_transformer.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/can/pre_gmm_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/can/pre_gmm_transformer.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/lift/pre_diffusion_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/lift/pre_diffusion_mlp_img.yaml

@@ -42,7 +42,6 @@ train:
     first_cycle_steps: 8000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/lift/pre_diffusion_unet.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/lift/pre_gaussian_mlp.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/lift/pre_gaussian_mlp_img.yaml

@@ -41,7 +41,6 @@ train:
     first_cycle_steps: 3000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/lift/pre_gaussian_transformer.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/lift/pre_gmm_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/lift/pre_gmm_transformer.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_diffusion_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_diffusion_mlp_img.yaml

@@ -42,7 +42,6 @@ train:
     first_cycle_steps: 8000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_diffusion_mlp_ta1.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_diffusion_mlp_ta1_ph.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_diffusion_unet.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_gaussian_mlp.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_gaussian_mlp_ibrl.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-4
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_gaussian_mlp_img.yaml

@@ -41,7 +41,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_gaussian_mlp_ta1_ph.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_gaussian_transformer.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_gmm_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/square/pre_gmm_transformer.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/transport/pre_diffusion_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/transport/pre_diffusion_mlp_img.yaml

@@ -42,7 +42,6 @@ train:
     first_cycle_steps: 8000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/transport/pre_diffusion_unet.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 10000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/transport/pre_gaussian_mlp.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/transport/pre_gaussian_mlp_img.yaml

@@ -41,7 +41,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/transport/pre_gaussian_transformer.yaml

@@ -31,7 +31,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/transport/pre_gmm_mlp.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model:

cfg/robomimic/pretrain/transport/pre_gmm_transformer.yaml

@@ -32,7 +32,6 @@ train:
     first_cycle_steps: 5000
     warmup_steps: 100
     min_lr: 1e-5
-  epoch_start_ema: 20
   save_model_freq: 500
 
 model: