Issue of action heads for version 1.5

[BUAAZhangHaonan]

Thank you very much for your update and brings us new things!
The performance after the update has indeed improved, but there are still some problems. I have some questions about the new action heads:

1.Problem of DiffusionActionHead

The updated DiffusionActionHead seems to be more unstable than before. After 50,000 steps of fine-tuning in the ALOHA environment, it still cannot be simulated normally, and an error of the MuJoCo will appear.

2.Problem of UNetDDPMActionHead

Thank you for providing the new UNetDDPMActionHead! But there seems to be some problems: In the ConditionalUnet1D of unet.py, the Downsample1d and Upsample1d will cause changes in dimensions, especially when the dimensions are odd, which may lead to wrong results. For example, when the dimension of action is (1, 1, 25, 512), Downsample1d will change it to (1, 1, 13, 1024), but in the Upsample1d, it will become (1, 1, 26, 512), which cannot match the hidden_reps dimension (1, 1, 25, 512). Changing action_horizon to 64 can solve this problem, but it cannot inference correctly. The following are my settings when finetuning:

config["model"]["heads"]["action"] = ModuleSpec.create(
UNetDDPMActionHead,
use_map=True,
flatten_tokens=False,
action_horizon=64,
action_dim=14,
readout_key="readout_action",
)