Set up proper models in eval

README.md

@@ -138,7 +138,7 @@ See [here](cfg/finetuning.md) for details of the experiments in the paper.
 * Videos of trials in Robomimic tasks can be recorded by specifying `env.save_video=True`, `train.render.freq=<iterations>`, and `train.render.num=<num_video>` in fine-tuning configs.
 
 ## Usage - Evaluation
-Pre-trained or fine-tuned policies can be evaluated without running the fine-tuning script now. Some example configs are provided under `cfg/{gym/robomimic/furniture}/eval}` including ones below. Set `base_policy_path` to override the default checkpoint. 
+Pre-trained or fine-tuned policies can be evaluated without running the fine-tuning script now. Some example configs are provided under `cfg/{gym/robomimic/furniture}/eval}` including ones below. `ft_denoising_steps` needs to match fine-tuning config. Set `base_policy_path` to override the default checkpoint. 
 ```console
 python script/run.py --config-name=eval_diffusion_mlp \
     --config-dir=cfg/gym/eval/hopper-v2

agent/finetune/train_agent.py

@@ -128,7 +128,7 @@ def save_model(self):
         data = {
             "itr": self.itr,
             "model": self.model.state_dict(),
-        }
+        }  # right now `model` includes weights for `network`, `actor`, `actor_ft`. Weights for `network` is redundant, and we can use `actor` weights as the base policy (earlier denoising steps) and `actor_ft` weights as the fine-tuned policy (later denoising steps) during evaluation.
         savepath = os.path.join(self.checkpoint_dir, f"state_{self.itr}.pt")
         torch.save(data, savepath)
         log.info(f"Saved model to {savepath}")

cfg/d3il/eval/avoid_m1/eval_diffusion_mlp.yaml

@@ -19,6 +19,7 @@ denoising_steps: 20
 cond_steps: 1
 horizon_steps: 4
 act_steps: 4
+ft_denoising_steps: 10
 
 n_steps: 25
 render_num: 40
@@ -47,7 +48,8 @@ env:
       reset_within_step: False
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   #

cfg/furniture/eval/lamp_low/eval_diffusion_mlp.yaml

@@ -21,6 +21,7 @@ horizon_steps: 8
 act_steps: 8
 use_ddim: True
 ddim_steps: 5
+ft_denoising_steps: 5
 
 n_steps: ${eval:'round(${env.max_episode_steps} / ${act_steps})'}
 render_num: 0
@@ -41,7 +42,8 @@ env:
     sparse_reward: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/furniture/eval/lamp_low/eval_diffusion_unet.yaml

@@ -21,6 +21,7 @@ horizon_steps: 16
 act_steps: 8
 use_ddim: True
 ddim_steps: 5
+ft_denoising_steps: 5
 
 n_steps: ${eval:'round(${env.max_episode_steps} / ${act_steps})'}
 render_num: 0
@@ -41,7 +42,8 @@ env:
     sparse_reward: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/furniture/eval/one_leg_low/eval_diffusion_mlp.yaml

@@ -21,6 +21,7 @@ horizon_steps: 8
 act_steps: 8
 use_ddim: True
 ddim_steps: 5
+ft_denoising_steps: 5
 
 n_steps: ${eval:'round(${env.max_episode_steps} / ${act_steps})'}
 render_num: 0
@@ -41,7 +42,8 @@ env:
     sparse_reward: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/furniture/eval/one_leg_low/eval_diffusion_unet.yaml

@@ -21,6 +21,7 @@ horizon_steps: 16
 act_steps: 8
 use_ddim: True
 ddim_steps: 5
+ft_denoising_steps: 5
 
 n_steps: ${eval:'round(${env.max_episode_steps} / ${act_steps})'}
 render_num: 0
@@ -41,7 +42,8 @@ env:
     sparse_reward: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/furniture/eval/round_table_low/eval_diffusion_mlp.yaml

@@ -21,6 +21,7 @@ horizon_steps: 8
 act_steps: 8
 use_ddim: True
 ddim_steps: 5
+ft_denoising_steps: 5
 
 n_steps: ${eval:'round(${env.max_episode_steps} / ${act_steps})'}
 render_num: 0
@@ -41,7 +42,8 @@ env:
     sparse_reward: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/furniture/eval/round_table_low/eval_diffusion_unet.yaml

@@ -21,6 +21,7 @@ horizon_steps: 16
 act_steps: 8
 use_ddim: True
 ddim_steps: 5
+ft_denoising_steps: 5
 
 n_steps: ${eval:'round(${env.max_episode_steps} / ${act_steps})'}
 render_num: 0
@@ -41,7 +42,8 @@ env:
     sparse_reward: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/gym/eval/halfcheetah-v2/eval_diffusion_mlp.yaml

@@ -19,6 +19,7 @@ denoising_steps: 20
 cond_steps: 1
 horizon_steps: 4
 act_steps: 4
+ft_denoising_steps: 10
 
 n_steps: 250  # each episode can take maximum (max_episode_steps / act_steps, =250 right now) steps but may finish earlier in gym. We only count episodes finished within n_steps for evaluation.
 render_num: 0
@@ -40,7 +41,8 @@ env:
       reset_within_step: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   #

cfg/gym/eval/hopper-v2/eval_diffusion_mlp.yaml

@@ -19,6 +19,7 @@ denoising_steps: 20
 cond_steps: 1
 horizon_steps: 4
 act_steps: 4
+ft_denoising_steps: 10
 
 n_steps: 250  # each episode can take maximum (max_episode_steps / act_steps, =250 right now) steps but may finish earlier in gym. We only count episodes finished within n_steps for evaluation.
 render_num: 0
@@ -40,7 +41,8 @@ env:
       reset_within_step: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   #

cfg/gym/eval/kitchen-v0/eval_diffusion_mlp.yaml

@@ -19,6 +19,7 @@ denoising_steps: 20
 cond_steps: 1
 horizon_steps: 4
 act_steps: 4
+ft_denoising_steps: 10
 
 n_steps: 70
 render_num: 0
@@ -40,7 +41,8 @@ env:
       reset_within_step: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   #

cfg/gym/eval/walker2d-v2/eval_diffusion_mlp.yaml

@@ -19,6 +19,7 @@ denoising_steps: 20
 cond_steps: 1
 horizon_steps: 4
 act_steps: 4
+ft_denoising_steps: 10
 
 n_steps: 250  # each episode can take maximum (max_episode_steps / act_steps, =250 right now) steps but may finish earlier in gym. We only count episodes finished within n_steps for evaluation.
 render_num: 0
@@ -40,7 +41,8 @@ env:
       reset_within_step: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   #

cfg/robomimic/eval/can/eval_diffusion_mlp.yaml

@@ -20,6 +20,7 @@ denoising_steps: 20
 cond_steps: 1
 horizon_steps: 4
 act_steps: 4
+ft_denoising_steps: 10
 
 n_steps: 300  # each episode takes max_episode_steps / act_steps steps
 render_num: 0
@@ -44,7 +45,8 @@ env:
       reset_within_step: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/robomimic/eval/can/eval_diffusion_mlp_img.yaml

@@ -23,6 +23,7 @@ horizon_steps: 4
 act_steps: 4
 use_ddim: True
 ddim_steps: 5
+ft_denoising_steps: 5
 
 n_steps: 300  # each episode takes max_episode_steps / act_steps steps
 render_num: 0
@@ -58,7 +59,8 @@ shape_meta:
     shape: [7]
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/robomimic/eval/can/eval_diffusion_unet.yaml

@@ -20,6 +20,7 @@ denoising_steps: 20
 cond_steps: 1
 horizon_steps: 4
 act_steps: 4
+ft_denoising_steps: 10
 
 n_steps: 75  # each episode takes max_episode_steps / act_steps steps
 render_num: 0
@@ -44,7 +45,8 @@ env:
       reset_within_step: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/robomimic/eval/can/eval_diffusion_unet_img.yaml

@@ -23,6 +23,7 @@ horizon_steps: 4
 act_steps: 4
 use_ddim: True
 ddim_steps: 5
+ft_denoising_steps: 5
 
 n_steps: 300  # each episode takes max_episode_steps / act_steps steps
 render_num: 0
@@ -58,7 +59,8 @@ shape_meta:
     shape: [7]
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/robomimic/eval/lift/eval_diffusion_mlp.yaml

@@ -20,6 +20,7 @@ denoising_steps: 20
 cond_steps: 1
 horizon_steps: 4
 act_steps: 4
+ft_denoising_steps: 10
 
 n_steps: 300  # each episode takes max_episode_steps / act_steps steps
 render_num: 0
@@ -44,7 +45,8 @@ env:
       reset_within_step: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/robomimic/eval/lift/eval_diffusion_mlp_img.yaml

@@ -23,6 +23,7 @@ horizon_steps: 4
 act_steps: 4
 use_ddim: True
 ddim_steps: 5
+ft_denoising_steps: 5
 
 n_steps: 300  # each episode takes max_episode_steps / act_steps steps
 render_num: 0
@@ -58,7 +59,8 @@ shape_meta:
     shape: [7]
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/robomimic/eval/lift/eval_diffusion_unet.yaml

@@ -20,6 +20,7 @@ denoising_steps: 20
 cond_steps: 1
 horizon_steps: 4
 act_steps: 4
+ft_denoising_steps: 10
 
 n_steps: 75  # each episode takes max_episode_steps / act_steps steps
 render_num: 0
@@ -44,7 +45,8 @@ env:
       reset_within_step: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/robomimic/eval/lift/eval_diffusion_unet_img.yaml

@@ -23,6 +23,7 @@ horizon_steps: 4
 act_steps: 4
 use_ddim: True
 ddim_steps: 5
+ft_denoising_steps: 5
 
 n_steps: 300  # each episode takes max_episode_steps / act_steps steps
 render_num: 0
@@ -58,7 +59,8 @@ shape_meta:
     shape: [7]
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/robomimic/eval/square/eval_diffusion_mlp.yaml

@@ -20,6 +20,7 @@ denoising_steps: 20
 cond_steps: 1
 horizon_steps: 4
 act_steps: 4
+ft_denoising_steps: 10
 
 n_steps: 400  # each episode takes max_episode_steps / act_steps steps
 render_num: 0
@@ -44,7 +45,8 @@ env:
       reset_within_step: True
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3

cfg/robomimic/eval/square/eval_diffusion_mlp_img.yaml

@@ -23,6 +23,7 @@ horizon_steps: 4
 act_steps: 4
 use_ddim: True
 ddim_steps: 5
+ft_denoising_steps: 5
 
 n_steps: 400  # each episode takes max_episode_steps / act_steps steps
 render_num: 0
@@ -58,7 +59,8 @@ shape_meta:
     shape: [7]
 
 model:
-  _target_: model.diffusion.diffusion.DiffusionModel
+  _target_: model.diffusion.diffusion_eval_ft.DiffusionEvalFT
+  ft_denoising_steps: ${ft_denoising_steps}
   predict_epsilon: True
   denoised_clip_value: 1.0
   randn_clip_value: 3