diff --git a/agent/pretrain/train_diffusion_agent.py b/agent/pretrain/train_diffusion_agent.py index e7747b2..44ed7fe 100644 --- 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() diff --git a/agent/pretrain/train_gaussian_agent.py b/agent/pretrain/train_gaussian_agent.py index 7384785..f02b03e 100644 --- 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() diff --git a/cfg/d3il/pretrain/avoid_m1/pre_diffusion_mlp.yaml b/cfg/d3il/pretrain/avoid_m1/pre_diffusion_mlp.yaml index e9e46bd..d3148f7 100644 --- 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 diff --git a/cfg/d3il/pretrain/avoid_m1/pre_gaussian_mlp.yaml b/cfg/d3il/pretrain/avoid_m1/pre_gaussian_mlp.yaml index ad49b8f..ae26e2b 100644 --- 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 diff --git a/cfg/d3il/pretrain/avoid_m1/pre_gmm_mlp.yaml b/cfg/d3il/pretrain/avoid_m1/pre_gmm_mlp.yaml index eae560a..16d9106 100644 --- 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 diff --git a/cfg/d3il/pretrain/avoid_m2/pre_diffusion_mlp.yaml b/cfg/d3il/pretrain/avoid_m2/pre_diffusion_mlp.yaml index 80c6c81..e49c7ec 100644 --- 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 diff --git a/cfg/d3il/pretrain/avoid_m2/pre_gaussian_mlp.yaml b/cfg/d3il/pretrain/avoid_m2/pre_gaussian_mlp.yaml index fc93398..dfbe78b 100644 --- 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 diff --git a/cfg/d3il/pretrain/avoid_m2/pre_gmm_mlp.yaml b/cfg/d3il/pretrain/avoid_m2/pre_gmm_mlp.yaml index 792fab5..2dc8999 100644 --- 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 diff --git a/cfg/d3il/pretrain/avoid_m3/pre_diffusion_mlp.yaml b/cfg/d3il/pretrain/avoid_m3/pre_diffusion_mlp.yaml index 292493c..c160859 100644 --- 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 diff --git a/cfg/d3il/pretrain/avoid_m3/pre_gaussian_mlp.yaml b/cfg/d3il/pretrain/avoid_m3/pre_gaussian_mlp.yaml index 587d0ba..cd0c24a 100644 --- 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 diff --git a/cfg/d3il/pretrain/avoid_m3/pre_gmm_mlp.yaml b/cfg/d3il/pretrain/avoid_m3/pre_gmm_mlp.yaml index 07f5a0c..82f8c35 100644 --- 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 diff --git a/cfg/furniture/pretrain/lamp_low/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/lamp_low/pre_diffusion_mlp.yaml index 4b655d4..706e099 100644 --- 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 diff --git a/cfg/furniture/pretrain/lamp_low/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/lamp_low/pre_diffusion_unet.yaml index 45b0d2e..1f09f63 100644 --- 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 diff --git a/cfg/furniture/pretrain/lamp_low/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/lamp_low/pre_gaussian_mlp.yaml index bc451c7..5cb32db 100644 --- 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 diff --git a/cfg/furniture/pretrain/lamp_med/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/lamp_med/pre_diffusion_mlp.yaml index d24459b..cedcc60 100644 --- 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 diff --git a/cfg/furniture/pretrain/lamp_med/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/lamp_med/pre_diffusion_unet.yaml index 478562f..1b30fdb 100644 --- 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 diff --git a/cfg/furniture/pretrain/lamp_med/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/lamp_med/pre_gaussian_mlp.yaml index e74dcd4..26361cf 100644 --- 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 diff --git a/cfg/furniture/pretrain/one_leg_low/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/one_leg_low/pre_diffusion_mlp.yaml index e751175..40c5e77 100644 --- 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 diff --git a/cfg/furniture/pretrain/one_leg_low/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/one_leg_low/pre_diffusion_unet.yaml index e981119..672f59f 100644 --- 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 diff --git a/cfg/furniture/pretrain/one_leg_low/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/one_leg_low/pre_gaussian_mlp.yaml index ff9919c..6ce994e 100644 --- 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 diff --git a/cfg/furniture/pretrain/one_leg_med/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/one_leg_med/pre_diffusion_mlp.yaml index 1ce6cd9..18786c6 100644 --- 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 diff --git a/cfg/furniture/pretrain/one_leg_med/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/one_leg_med/pre_diffusion_unet.yaml index bca3be5..1d90175 100644 --- 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 diff --git a/cfg/furniture/pretrain/one_leg_med/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/one_leg_med/pre_gaussian_mlp.yaml index 4cce367..ce0066a 100644 --- 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 diff --git a/cfg/furniture/pretrain/round_table_low/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/round_table_low/pre_diffusion_mlp.yaml index 5ac604e..7a56d78 100644 --- 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 diff --git a/cfg/furniture/pretrain/round_table_low/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/round_table_low/pre_diffusion_unet.yaml index 9d20e49..32f2127 100644 --- 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 diff --git a/cfg/furniture/pretrain/round_table_low/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/round_table_low/pre_gaussian_mlp.yaml index 218559e..46b1c39 100644 --- 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 diff --git a/cfg/furniture/pretrain/round_table_med/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/round_table_med/pre_diffusion_mlp.yaml index d4e0c56..074f9ea 100644 --- 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 diff --git a/cfg/furniture/pretrain/round_table_med/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/round_table_med/pre_diffusion_unet.yaml index 1965d40..73da6cc 100644 --- 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 diff --git a/cfg/furniture/pretrain/round_table_med/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/round_table_med/pre_gaussian_mlp.yaml index 1ade284..a29708c 100644 --- 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 diff --git a/cfg/gym/pretrain/halfcheetah-medium-v2/pre_diffusion_mlp.yaml b/cfg/gym/pretrain/halfcheetah-medium-v2/pre_diffusion_mlp.yaml index 612bc5e..d17ccad 100644 --- 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: diff --git a/cfg/gym/pretrain/halfcheetah-medium-v2/pre_gaussian_mlp.yaml b/cfg/gym/pretrain/halfcheetah-medium-v2/pre_gaussian_mlp.yaml index 9a7f5bf..2af034e 100644 --- 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: diff --git a/cfg/gym/pretrain/hopper-medium-v2/pre_diffusion_mlp.yaml b/cfg/gym/pretrain/hopper-medium-v2/pre_diffusion_mlp.yaml index 4b05bba..5dc8f75 100644 --- 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: diff --git a/cfg/gym/pretrain/hopper-medium-v2/pre_gaussian_mlp.yaml b/cfg/gym/pretrain/hopper-medium-v2/pre_gaussian_mlp.yaml index 5a11734..3ec109b 100644 --- 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: diff --git a/cfg/gym/pretrain/kitchen-complete-v0/pre_diffusion_mlp.yaml b/cfg/gym/pretrain/kitchen-complete-v0/pre_diffusion_mlp.yaml index 092fa00..b0f6a85 100644 --- 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 diff --git a/cfg/gym/pretrain/kitchen-complete-v0/pre_gaussian_mlp.yaml b/cfg/gym/pretrain/kitchen-complete-v0/pre_gaussian_mlp.yaml index 9426b6f..a8f72d2 100644 --- 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 diff --git a/cfg/gym/pretrain/kitchen-mixed-v0/pre_diffusion_mlp.yaml b/cfg/gym/pretrain/kitchen-mixed-v0/pre_diffusion_mlp.yaml index becf244..677143d 100644 --- 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 diff --git a/cfg/gym/pretrain/kitchen-mixed-v0/pre_gaussian_mlp.yaml b/cfg/gym/pretrain/kitchen-mixed-v0/pre_gaussian_mlp.yaml index 86a6e90..d0a64e9 100644 --- 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 diff --git a/cfg/gym/pretrain/kitchen-partial-v0/pre_diffusion_mlp.yaml b/cfg/gym/pretrain/kitchen-partial-v0/pre_diffusion_mlp.yaml index c854707..95e7d16 100644 --- 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 diff --git a/cfg/gym/pretrain/kitchen-partial-v0/pre_gaussian_mlp.yaml b/cfg/gym/pretrain/kitchen-partial-v0/pre_gaussian_mlp.yaml index 02413a5..0796ff3 100644 --- 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 diff --git a/cfg/gym/pretrain/walker2d-medium-v2/pre_diffusion_mlp.yaml b/cfg/gym/pretrain/walker2d-medium-v2/pre_diffusion_mlp.yaml index 8118aba..4ed8dd8 100644 --- 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: diff --git a/cfg/gym/pretrain/walker2d-medium-v2/pre_gaussian_mlp.yaml b/cfg/gym/pretrain/walker2d-medium-v2/pre_gaussian_mlp.yaml index dca4923..346d068 100644 --- 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: diff --git a/cfg/pretraining.md b/cfg/pretraining.md index c672d17..84a2b5a 100644 --- 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//`, and the config name is `pre_diffusion_mlp`. Robomimic configs are under `cfg/robomimic/pretrain//`, and the name is also `pre_diffusion_mlp`. diff --git a/cfg/robomimic/finetune/can/ibrl_mlp.yaml b/cfg/robomimic/finetune/can/ibrl_mlp.yaml index c3ba56e..ea9dae6 100644 --- 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 diff --git a/cfg/robomimic/finetune/can/ibrl_mlp_ph.yaml b/cfg/robomimic/finetune/can/ibrl_mlp_ph.yaml index 8940658..8741102 100644 --- 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 diff --git a/cfg/robomimic/finetune/square/ibrl_mlp.yaml b/cfg/robomimic/finetune/square/ibrl_mlp.yaml index fba5969..f399c71 100644 --- 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 diff --git a/cfg/robomimic/finetune/square/ibrl_mlp_ph.yaml b/cfg/robomimic/finetune/square/ibrl_mlp_ph.yaml index f65c1dd..a869467 100644 --- 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 diff --git a/cfg/robomimic/pretrain/can/pre_diffusion_mlp.yaml b/cfg/robomimic/pretrain/can/pre_diffusion_mlp.yaml index 8834aee..038182d 100644 --- 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 diff --git a/cfg/robomimic/pretrain/can/pre_diffusion_mlp_img.yaml b/cfg/robomimic/pretrain/can/pre_diffusion_mlp_img.yaml index 2ce38ef..b09e0ee 100644 --- 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: diff --git a/cfg/robomimic/pretrain/can/pre_diffusion_mlp_ta1.yaml b/cfg/robomimic/pretrain/can/pre_diffusion_mlp_ta1.yaml index 62a09e8..fb5a9cb 100644 --- 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 diff --git a/cfg/robomimic/pretrain/can/pre_diffusion_mlp_ta1_ph.yaml b/cfg/robomimic/pretrain/can/pre_diffusion_mlp_ta1_ph.yaml index 46593c6..1fb3656 100644 --- 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 diff --git a/cfg/robomimic/pretrain/can/pre_diffusion_unet.yaml b/cfg/robomimic/pretrain/can/pre_diffusion_unet.yaml index 14aac61..8284135 100644 --- 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 diff --git a/cfg/robomimic/pretrain/can/pre_gaussian_mlp.yaml b/cfg/robomimic/pretrain/can/pre_gaussian_mlp.yaml index 97800c2..9840bad 100644 --- 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 diff --git a/cfg/robomimic/pretrain/can/pre_gaussian_mlp_ibrl.yaml b/cfg/robomimic/pretrain/can/pre_gaussian_mlp_ibrl.yaml index f07ffd8..272c173 100644 --- 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 diff --git a/cfg/robomimic/pretrain/can/pre_gaussian_mlp_img.yaml b/cfg/robomimic/pretrain/can/pre_gaussian_mlp_img.yaml index f68d400..3f5cdf5 100644 --- 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: diff --git a/cfg/robomimic/pretrain/can/pre_gaussian_mlp_ta1_ph.yaml b/cfg/robomimic/pretrain/can/pre_gaussian_mlp_ta1_ph.yaml index 1bb170d..88d330e 100644 --- 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 diff --git a/cfg/robomimic/pretrain/can/pre_gaussian_transformer.yaml b/cfg/robomimic/pretrain/can/pre_gaussian_transformer.yaml index 37e9fcf..9a737db 100644 --- 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 diff --git a/cfg/robomimic/pretrain/can/pre_gmm_mlp.yaml b/cfg/robomimic/pretrain/can/pre_gmm_mlp.yaml index ebef165..8716bf6 100644 --- 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 diff --git a/cfg/robomimic/pretrain/can/pre_gmm_transformer.yaml b/cfg/robomimic/pretrain/can/pre_gmm_transformer.yaml index 3c698aa..a5e3524 100644 --- 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 diff --git a/cfg/robomimic/pretrain/lift/pre_diffusion_mlp.yaml b/cfg/robomimic/pretrain/lift/pre_diffusion_mlp.yaml index 94c31d9..1555fa5 100644 --- 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 diff --git a/cfg/robomimic/pretrain/lift/pre_diffusion_mlp_img.yaml b/cfg/robomimic/pretrain/lift/pre_diffusion_mlp_img.yaml index 418daac..3f5ddc5 100644 --- 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: diff --git a/cfg/robomimic/pretrain/lift/pre_diffusion_unet.yaml b/cfg/robomimic/pretrain/lift/pre_diffusion_unet.yaml index cb85f89..07a7ff1 100644 --- 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 diff --git a/cfg/robomimic/pretrain/lift/pre_gaussian_mlp.yaml b/cfg/robomimic/pretrain/lift/pre_gaussian_mlp.yaml index 616c27c..58285a7 100644 --- 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 diff --git a/cfg/robomimic/pretrain/lift/pre_gaussian_mlp_img.yaml b/cfg/robomimic/pretrain/lift/pre_gaussian_mlp_img.yaml index e31d30c..a3de5f3 100644 --- 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: diff --git a/cfg/robomimic/pretrain/lift/pre_gaussian_transformer.yaml b/cfg/robomimic/pretrain/lift/pre_gaussian_transformer.yaml index 50add2d..4067a53 100644 --- 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 diff --git a/cfg/robomimic/pretrain/lift/pre_gmm_mlp.yaml b/cfg/robomimic/pretrain/lift/pre_gmm_mlp.yaml index 24fd45a..0434a78 100644 --- 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 diff --git a/cfg/robomimic/pretrain/lift/pre_gmm_transformer.yaml b/cfg/robomimic/pretrain/lift/pre_gmm_transformer.yaml index f24767a..7bee492 100644 --- 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 diff --git a/cfg/robomimic/pretrain/square/pre_diffusion_mlp.yaml b/cfg/robomimic/pretrain/square/pre_diffusion_mlp.yaml index 4f4672b..2e15072 100644 --- 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 diff --git a/cfg/robomimic/pretrain/square/pre_diffusion_mlp_img.yaml b/cfg/robomimic/pretrain/square/pre_diffusion_mlp_img.yaml index 8da14e0..c81ba12 100644 --- 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: diff --git a/cfg/robomimic/pretrain/square/pre_diffusion_mlp_ta1.yaml b/cfg/robomimic/pretrain/square/pre_diffusion_mlp_ta1.yaml index 53e572e..9da7f6e 100644 --- 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 diff --git a/cfg/robomimic/pretrain/square/pre_diffusion_mlp_ta1_ph.yaml b/cfg/robomimic/pretrain/square/pre_diffusion_mlp_ta1_ph.yaml index 7bffecd..de9c232 100644 --- 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 diff --git a/cfg/robomimic/pretrain/square/pre_diffusion_unet.yaml b/cfg/robomimic/pretrain/square/pre_diffusion_unet.yaml index b5a8faf..6811bab 100644 --- 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 diff --git a/cfg/robomimic/pretrain/square/pre_gaussian_mlp.yaml b/cfg/robomimic/pretrain/square/pre_gaussian_mlp.yaml index 652c7bd..19852a9 100644 --- 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 diff --git a/cfg/robomimic/pretrain/square/pre_gaussian_mlp_ibrl.yaml b/cfg/robomimic/pretrain/square/pre_gaussian_mlp_ibrl.yaml index 38f3402..992bb99 100644 --- 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 diff --git a/cfg/robomimic/pretrain/square/pre_gaussian_mlp_img.yaml b/cfg/robomimic/pretrain/square/pre_gaussian_mlp_img.yaml index 7bb5632..8728178 100644 --- 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: diff --git a/cfg/robomimic/pretrain/square/pre_gaussian_mlp_ta1_ph.yaml b/cfg/robomimic/pretrain/square/pre_gaussian_mlp_ta1_ph.yaml index 0cc2204..e44b09c 100644 --- 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 diff --git a/cfg/robomimic/pretrain/square/pre_gaussian_transformer.yaml b/cfg/robomimic/pretrain/square/pre_gaussian_transformer.yaml index a469afa..320637f 100644 --- 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 diff --git a/cfg/robomimic/pretrain/square/pre_gmm_mlp.yaml b/cfg/robomimic/pretrain/square/pre_gmm_mlp.yaml index 315c980..95c9273 100644 --- 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 diff --git a/cfg/robomimic/pretrain/square/pre_gmm_transformer.yaml b/cfg/robomimic/pretrain/square/pre_gmm_transformer.yaml index 27afc79..61253fc 100644 --- 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 diff --git a/cfg/robomimic/pretrain/transport/pre_diffusion_mlp.yaml b/cfg/robomimic/pretrain/transport/pre_diffusion_mlp.yaml index 2c6f868..053797d 100644 --- 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 diff --git a/cfg/robomimic/pretrain/transport/pre_diffusion_mlp_img.yaml b/cfg/robomimic/pretrain/transport/pre_diffusion_mlp_img.yaml index 54415dc..c9d30ed 100644 --- 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: diff --git a/cfg/robomimic/pretrain/transport/pre_diffusion_unet.yaml b/cfg/robomimic/pretrain/transport/pre_diffusion_unet.yaml index 15dd4d9..71bd096 100644 --- 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 diff --git a/cfg/robomimic/pretrain/transport/pre_gaussian_mlp.yaml b/cfg/robomimic/pretrain/transport/pre_gaussian_mlp.yaml index 08cb96f..3d88907 100644 --- 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 diff --git a/cfg/robomimic/pretrain/transport/pre_gaussian_mlp_img.yaml b/cfg/robomimic/pretrain/transport/pre_gaussian_mlp_img.yaml index 1ce968b..5d66544 100644 --- 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: diff --git a/cfg/robomimic/pretrain/transport/pre_gaussian_transformer.yaml b/cfg/robomimic/pretrain/transport/pre_gaussian_transformer.yaml index ffb4be8..7c81253 100644 --- 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 diff --git a/cfg/robomimic/pretrain/transport/pre_gmm_mlp.yaml b/cfg/robomimic/pretrain/transport/pre_gmm_mlp.yaml index dbe6290..9879fee 100644 --- 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 diff --git a/cfg/robomimic/pretrain/transport/pre_gmm_transformer.yaml b/cfg/robomimic/pretrain/transport/pre_gmm_transformer.yaml index 17525d5..c2a0ebe 100644 --- 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 diff --git a/cfg/robomimic/scratch/can/rlpd_mlp.yaml b/cfg/robomimic/scratch/can/rlpd_mlp.yaml index 8b66075..a62c9ef 100644 --- 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 diff --git a/cfg/robomimic/scratch/can/rlpd_mlp_ph.yaml b/cfg/robomimic/scratch/can/rlpd_mlp_ph.yaml index d574d5a..8fde067 100644 --- 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 diff --git a/cfg/robomimic/scratch/square/rlpd_mlp.yaml b/cfg/robomimic/scratch/square/rlpd_mlp.yaml index 46730a7..900efbf 100644 --- 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 diff --git a/cfg/robomimic/scratch/square/rlpd_mlp_ph.yaml b/cfg/robomimic/scratch/square/rlpd_mlp_ph.yaml index cb8a8b4..687c273 100644 --- 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 diff --git a/script/download_url.py b/script/download_url.py index 1fa2069..d50ee4d 100644 --- 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