GDDPM implementation

.gitignore

@@ -88,10 +88,16 @@ dmypy.json
 multirun/
 
 # Model weights
-weights/
+weights/*
 
 # Archive files
 archive/
 
 # Notebooks
-notebooks/
+notebooks/
+
+# Docs
+docs/
+
+# Claude
+CLAUDE.md

conda_environment.yaml

@@ -14,8 +14,8 @@ dependencies:
     - pytorchvideo
     - gymnasium==0.28.1
     - gym
-    - -e git+https://github.com/ARISE-Initiative/robomimic@main#egg=robomimic
-    - diffusers
+    - robomimic
+    - diffusers==0.29.2
     - zarr
     - einops
     - tqdm
@@ -33,6 +33,5 @@ dependencies:
     - wandb
     # related work dependencies
     - -e git+https://github.com/columbia-ai-robotics/diffusion_policy@main#egg=diffusion_policy
-    - -e git+https://github.com/anindex/stoch_gpmp@main#egg=stoch_gpmp
     - -e git+https://github.com/anindex/torch_robotics@main#egg=torch_robotics
 

imitation/config/policy/argd_policy.yaml

@@ -16,4 +16,5 @@ denoising_network:
   num_edge_types: ${policy.num_edge_types}
   num_layers: 2
   hidden_dim: 512
-use_normalization: False
+use_normalization: False
+num_warmup_steps: 100

imitation/config/policy/diffusion_policy.yaml

@@ -7,4 +7,5 @@ action_horizon: ${action_horizon}
 num_diffusion_iters: 100
 dataset: ${task.dataset}
 ckpt_path: ./weights/diffusion_policy_${task.task_name}_${task.dataset_type}.pt
-lr: 0.00001
+lr: 0.00001
+num_warmup_steps: 100

imitation/config/policy/gddpm_policy.yaml

@@ -0,0 +1,35 @@
+_target_: imitation.policy.graph_ddpm_policy.GraphConditionalDDPMPolicy
+
+obs_dim: ${task.obs_dim}
+action_dim: ${task.action_dim}
+
+node_feature_dim: 1  # [joint_val]  (same as graph_ddpm_policy)
+num_edge_types: 2    # robot joints, object-robot
+pred_horizon: ${pred_horizon}
+obs_horizon: ${obs_horizon}
+action_horizon: ${action_horizon}
+num_diffusion_iters: 100
+dataset: ${task.dataset}
+
+denoising_network:
+  _target_: imitation.model.gddpm.GDDPMNoisePred
+  node_feature_dim: ${policy.node_feature_dim}
+  cond_feature_dim: 9        # 9-D obs features: pos(3) + rot_6d(6) (excl. node-id)
+  obs_horizon: ${obs_horizon}
+  pred_horizon: ${pred_horizon}
+  edge_feature_dim: 1
+  num_edge_types: ${policy.num_edge_types}
+  # GDDPM-specific hyper-parameters
+  residual_layers: 8
+  residual_channels: 32
+  dilation_cycle_length: 2
+  hidden_dim: 256
+  diffusion_step_embed_dim: 64
+  num_diffusion_steps: ${policy.num_diffusion_iters}
+
+ckpt_path: ./weights/gddpm_policy_${task.task_name}_${task.dataset_type}_${task.control_mode}_${policy.num_diffusion_iters}iters.pt
+lr: 1e-4
+batch_size: 128
+use_normalization: True
+keep_first_action: True
+num_warmup_steps: 100

imitation/config/policy/graph_ddpm_policy.yaml

@@ -23,7 +23,7 @@ denoising_network:
   diffusion_step_embed_dim: 64
   num_diffusion_steps: ${policy.num_diffusion_iters}
 ckpt_path: ./weights/diffusion_graph_policy_${task.task_name}_${task.dataset_type}_${task.control_mode}_${policy.num_diffusion_iters}iters.pt
-lr: 1e-4
+lr: 1e-5
 batch_size: 128
 use_normalization: True
 keep_first_action: True

imitation/config/policy/osc_ddpm_policy.yaml

@@ -0,0 +1,32 @@
+_target_: imitation.policy.osc_ddpm_policy.OSCGraphDDPMPolicy
+
+action_dim: ${task.action_dim}   # 7 for OSC_POSE
+num_edge_types: 2
+pred_horizon: ${pred_horizon}
+obs_horizon: ${obs_horizon}
+action_horizon: ${action_horizon}
+num_diffusion_iters: 50
+dataset: ${task.dataset}
+
+denoising_network:
+  _target_: imitation.model.gddpm.FlatGDDPMNoisePred
+  action_dim: ${policy.action_dim}
+  cond_feature_dim: 9        # 9-D obs features: pos(3) + rot_6d(6)
+  obs_horizon: ${obs_horizon}
+  pred_horizon: ${pred_horizon}
+  edge_feature_dim: 1
+  num_edge_types: ${policy.num_edge_types}
+  residual_layers: 8
+  residual_channels: 32
+  dilation_cycle_length: 2
+  hidden_dim: 256
+  diffusion_step_embed_dim: 64
+  num_diffusion_steps: ${policy.num_diffusion_iters}
+
+ckpt_path: ./weights/osc_ddpm_policy_${task.task_name}_${task.control_mode}_${policy.num_diffusion_iters}iters.pt
+lr: 1e-4
+batch_size: 128
+use_normalization: True
+keep_first_action: False
+num_warmup_steps: 100
+ema_decay: 0.9999

imitation/config/task/lift_graph_osc.yaml

@@ -0,0 +1,53 @@
+
+task_name: &task_name lift
+dataset_type: &dataset_type ph
+dataset_path: &dataset_path ./data/lift/${task.dataset_type}/low_dim_v141.hdf5
+
+max_steps: 500
+
+control_mode: "OSC_POSE"
+
+obs_dim: 10     # 9 robot + 1 object nodes
+action_dim: 7   # flat 7-D EEF vector (xyz + rotation + gripper)
+
+robots: ["Panda"]
+
+object_state_sizes: &object_state_sizes
+  cube_pos: 3
+  cube_quat: 4
+  gripper_to_cube_pos: 3
+
+object_state_keys: &object_state_keys
+  cube: ["cube_pos", "cube_quat"]
+
+env_runner:
+  _target_: imitation.env_runner.robomimic_lowdim_runner.RobomimicEnvRunner
+  output_dir: ${output_dir}
+  action_horizon: ${action_horizon}
+  obs_horizon: ${obs_horizon}
+  action_offset: ${action_offset}
+  render: ${render}
+  output_video: ${output_video}
+  use_full_pred_after: 0.4
+  env:
+    _target_: imitation.env.robomimic_graph_wrapper.RobomimicGraphWrapper
+    object_state_sizes: *object_state_sizes
+    object_state_keys:  *object_state_keys
+    max_steps: ${task.max_steps}
+    task: "Lift"
+    has_renderer: ${render}
+    robots: ${task.robots}
+    output_video: ${output_video}
+    control_mode: ${task.control_mode}
+    base_link_shift: [[-0.56, 0, 0.912]]
+
+dataset:
+  _target_: imitation.dataset.robomimic_graph_dataset.RobomimicGraphDataset
+  dataset_path: ${task.dataset_path}
+  robots: ${task.robots}
+  pred_horizon: ${pred_horizon}
+  obs_horizon: ${obs_horizon}
+  object_state_sizes: *object_state_sizes
+  object_state_keys:  *object_state_keys
+  control_mode: ${task.control_mode}
+  base_link_shift: [[-0.56, 0, 0.912]]

imitation/config/task/square_graph_osc.yaml

@@ -0,0 +1,55 @@
+
+task_name: &task_name square
+dataset_type: &dataset_type ph
+dataset_path: &dataset_path ./data/${task.task_name}/${task.dataset_type}/low_dim_v141.hdf5
+
+max_steps: ${eval:'800 if "${task.dataset_type}" == "mh" else 600'}
+
+control_mode: "OSC_POSE"
+
+obs_dim: 10     # 9 robot + 1 object nodes
+action_dim: 7   # flat 7-D EEF vector (xyz + rotation + gripper)
+
+robots: ["Panda"]
+
+object_state_sizes: &object_state_sizes
+  nut_pos: 3
+  nut_quat: 4
+
+object_state_keys: &object_state_keys
+  nut: ["nut_pos", "nut_quat"]
+
+env_runner:
+  _target_: imitation.env_runner.robomimic_lowdim_runner.RobomimicEnvRunner
+  output_dir: ${output_dir}
+  action_horizon: ${action_horizon}
+  obs_horizon: ${obs_horizon}
+  action_offset: ${action_offset}
+  render: ${render}
+  output_video: ${output_video}
+  use_full_pred_after: 0.8
+  env:
+    _target_: imitation.env.robomimic_graph_wrapper.RobomimicGraphWrapper
+    object_state_sizes: *object_state_sizes
+    object_state_keys:  *object_state_keys
+    max_steps: ${task.max_steps}
+    task: "NutAssemblySquare"
+    has_renderer: ${render}
+    robots: ${task.robots}
+    output_video: ${output_video}
+    control_mode: ${task.control_mode}
+    controller_config:
+      interpolation: "linear"
+      ramp_ratio: 0.2
+    base_link_shift: [[-0.56, 0, 0.912]]
+
+dataset:
+  _target_: imitation.dataset.robomimic_graph_dataset.RobomimicGraphDataset
+  dataset_path: ${task.dataset_path}
+  robots: ${task.robots}
+  pred_horizon: ${pred_horizon}
+  obs_horizon: ${obs_horizon}
+  object_state_sizes: *object_state_sizes
+  object_state_keys:  *object_state_keys
+  control_mode: ${task.control_mode}
+  base_link_shift: [[-0.56, 0, 0.912]]

imitation/config/test.yaml

@@ -1,7 +1,7 @@
 defaults:
   - _self_
   - task: lift_graph
-  - policy: graph_ddpm_policy
+  - policy: osc_ddpm_policy
 
 output_video: False
 render: True
@@ -13,8 +13,8 @@ output_dir: ./outputs
 
 pred_horizon: 16
 obs_horizon: 4
-action_horizon: 4
-action_offset: 1 # action offset for the policy, 1 if first action is to be ignored
+action_horizon: 12
+action_offset: 0 # action offset for the policy, 1 if first action is to be ignored
 
 
 agent:

imitation/config/train.yaml

@@ -1,7 +1,7 @@
 defaults:
   - _self_
   - task: lift_graph
-  - policy: graph_ddpm_policy
+  - policy: gddpm_policy
 
 output_dir: ./outputs
 # on evaluating, for environment wrapper
@@ -10,10 +10,10 @@ output_video: True
 
 pred_horizon: 16
 obs_horizon: 4
-action_horizon: 2
-action_offset: 1 # action offset for the policy, 1 if first action is to be ignored
+action_horizon: 4
+action_offset: 0 # action offset for the policy, 1 if first action is to be ignored
 # Training parameters
-num_epochs: 500
+num_epochs: 50
 val_fraction: 0.1
 seed: 0
 load_ckpt: False # start training from scratch
@@ -29,15 +29,15 @@ max_steps: ${task.max_steps}
 
 # Evaluation during training
 eval_params:
-  eval_every: 50 # evaluate every 50 epochs
+  eval_every: 25 # evaluate every 50 epochs
   val_every: 1
   task: ${task}
   policy: ${policy}
   render: False
   output_video: True
   load_ckpt: True # always load for evaluation
 
-  num_episodes: 50
+  num_episodes: 5
   max_steps: ${task.max_steps}
   output_dir: ./outputs