Error running robosuite Lift-ph with abs_action=false

Hi, thanks for the beautiful code. I am trying to run robosuite lift-ph with the non-abs-action dataset and abs_action=false. I downloaded the config from https://diffusion-policy.cs.columbia.edu/data/experiments/low_dim/lift_ph/diffusion_policy_cnn/config.yaml and renamed it as './configs/lowdim_liftph_diffusion_policy_cnn.yaml'. The codes run perfectly when I use the default settings where abs_action=true:

CUDA_VISIBLE_DEVICES=1 python train.py --config-dir=./configs --config-name=lowdim_liftph_diffusion_policy_cnn.yaml \
training.seed=42 training.device=cuda:0   \
task.abs_action=true  \
task.dataset.abs_action=true  task.dataset.dataset_path='data/robomimic/datasets/lift/ph/low_dim_abs.hdf5' \
task.dataset_path='data/robomimic/datasets/lift/ph/low_dim_abs.hdf5'  \
task.env_runner.abs_action=true   \
task.env_runner.dataset_path='data/robomimic/datasets/lift/ph/low_dim_abs.hdf5'     \
hydra.run.dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}'  \
multi_run.wandb_name_base='${now:%Y.%m.%d}_${now:%H.%M.%S}_${name}_${task_name}_${task.abs_action}_${task.env_runner.abs_action}'  \
logging.name='${now:%Y.%m.%d}_${now:%H.%M.%S}_${name}_${task_name}_${task.abs_action}_${task.env_runner.abs_action}'

However, an error occurred when I tried to set abs_action to false and the dataset to the non-abs-action version, with the command:

CUDA_VISIBLE_DEVICES=2 python train.py --config-dir=./configs --config-name=lowdim_liftph_diffusion_policy_cnn.yaml \
training.seed=42 training.device=cuda:0   \
task.abs_action=false  \
task.dataset.abs_action=false \
task.dataset.dataset_path='data/robomimic/datasets/lift/ph/low_dim.hdf5' \
task.dataset_path='data/robomimic/datasets/lift/ph/low_dim.hdf5'  \
task.env_runner.abs_action=false   \
task.env_runner.dataset_path='data/robomimic/datasets/lift/ph/low_dim.hdf5'   \
hydra.run.dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}' \
multi_run.wandb_name_base='${now:%Y.%m.%d}_${now:%H.%M.%S}_${name}_${task_name}_${task.abs_action}_${task.env_runner.abs_action}' \
logging.name='${now:%Y.%m.%d}_${now:%H.%M.%S}_${name}_${task_name}_${task.abs_action}_${task.env_runner.abs_action}'  \
logging.name='${now:%Y.%m.%d}_${now:%H.%M.%S}_${name}_${task_name}_${task.abs_action}_${task.env_runner.abs_action}'

It reports the error:

Error executing job with overrides: ['training.seed=42', 'training.device=cuda:0', 'task.abs_action=false', 'task.dataset.abs_action=false', 'task.dataset.dataset_
path=data/robomimic/datasets/lift/ph/low_dim.hdf5', 'task.dataset_path=data/robomimic/datasets/lift/ph/low_dim.hdf5', 'task.env_runner.abs_action=false', 'task.env
_runner.dataset_path=data/robomimic/datasets/lift/ph/low_dim.hdf5', 'multi_run.wandb_name_base=${now:%Y.%m.%d}_${now:%H.%M.%S}_${name}_${task_name}_${task.abs_acti
on}_${task.env_runner.abs_action}', 'logging.name=${now:%Y.%m.%d}_${now:%H.%M.%S}_${name}_${task_name}_${task.abs_action}_${task.env_runner.abs_action}', 'logging.
name=${now:%Y.%m.%d}_${now:%H.%M.%S}_${name}_${task_name}_${task.abs_action}_${task.env_runner.abs_action}']
Traceback (most recent call last):
  File "/data/zj/diffusion_policy/train.py", line 32, in main
    workspace.run()
  File "/data/zj/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_lowdim_workspace.py", line 168, in run
    raw_loss = self.model.compute_loss(batch)
  File "/data/zj/diffusion_policy/diffusion_policy/policy/diffusion_unet_lowdim_policy.py", line 215, in compute_loss
    condition_mask = self.mask_generator(trajectory.shape)
  File "/home/zj/anaconda3/envs/robodiff/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/zj/anaconda3/envs/robodiff/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/data/zj/diffusion_policy/diffusion_policy/model/diffusion/mask_generator.py", line 63, in forward
    assert D == (self.action_dim + self.obs_dim)
AssertionError

There seems to be a misalignment with the data and the env obs and action shapes. Do I miss any other configs that need to be changed in order to run the non-abs-action version? Or can you provide commands or configs for running these experiments? Thanks a lot.

real-stanford / diffusion_policy

Error running robosuite Lift-ph with abs_action=false #71