Actions are denoised twice in collect_rollouts of dsrl_sac

[fedeceola]

Hi @ajwagen, I am opening this issue because I believe that there is a bug in the collection of rollouts for dsrl_sac.

The noise is denoised twice by diffusion policy. The first denoising happens here:

dsrl/utils.py

Line 171 in 8d21b9c

action = base_policy(torch.tensor(obs, device=cfg.device, dtype=torch.float32), noise)

The second happens here:

dsrl/env_utils.py

Line 217 in 8d21b9c

diffused_actions = self.base_policy(self.obs, actions)

When the step function is called here:

dsrl/utils.py

Line 172 in 8d21b9c

next_obs, reward, done, info = env.step(action)

This issue is not affecting dsrl_na.

[omkarpatil18]

Another related but perhaps minor bug that also affects dsrl_na is here-

dsrl/train_dsrl.py

Line 173 in 8d21b9c

logging_callback.set_timesteps(cfg.train.init_rollout_steps * num_env)

Here the total timesteps should have been updated after multiplying with the chunk size, as done in the logger.

[fedeceola]

Thanks @omkarpatil18. If you are using dsrl_na I found another bug. It doesn't consider action_magnitude. As a consequence it's as if you always set action_magnitude equal to 1. You may want to fix it in these lines.

dsrl/env_utils.py

Lines 145 to 146 in 8d21b9c

	low=np.tile(env.action_space.low, cfg.act_steps),
	high=np.tile(env.action_space.high, cfg.act_steps),

[omkarpatil18]

Thanks @fedeceola! Another bug/improvement esp for DSRL SAC: Huge performance improvement on setting terminated and truncated correctly, instead of running all episodes to fixed length.

dsrl/env_utils.py

Line 184 in 8d21b9c

return obs, reward, done, False, info

With this fix there is little gain for DSRL NA over DSRL SAC.

[fedeceola]

Thanks! If you have already fixed it, would you be available to share the correct code?

[omkarpatil18]

Hi @fedeceola, apologies for the late response. Here is what I mean:

		terminated = check_success() or underlying_done
		truncated = False
		if self.count >= max_episode_steps:
			truncated = True
			info["TimeLimit.truncated"] = True
		if terminated or truncated:
			info["terminal_observation"] = obs
		return obs, reward, terminated, truncated, info

sets truncated and terminated correctly, instead of returning truncated=False done here:

dsrl/env_utils.py

Line 184 in 8d21b9c

return obs, reward, done, False, info

This may result in incorrect bootstrapping for false truncations (as in this case).