Shape mismatch when set clip_actions=True

[giangdao1402]

Hi @Toni-SM, thank for your great work !
I encountered a problem related to shape when trying to set clip_actions=True in both Policy and Value classes, when set both to False everything is ok
Here is the code that I currently use

class Policy(GaussianMixin, Model):
    def __init__(self, observation_space, action_space, device, clip_actions=True,
                 clip_log_std=True, min_log_std=-20, max_log_std=2, reduction="sum"):
        Model.__init__(self, observation_space, action_space, device)
        GaussianMixin.__init__(self, clip_actions, clip_log_std, min_log_std, max_log_std, reduction)

        self.features_extractor = nn.Sequential(
            nn.LazyConv2d(out_channels=32, kernel_size=8, stride=4, padding=0),
            nn.ReLU(),
            nn.LazyConv2d(out_channels=64, kernel_size=4, stride=2, padding=0),
            nn.ReLU(),
            nn.LazyConv2d(out_channels=64, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Flatten()
        )

        self.net_state = nn.Sequential(
            nn.LazyLinear(out_features=128),
            nn.ELU(),
            nn.LazyLinear(out_features=128),
            nn.ELU(),
            nn.LazyLinear(out_features=64),
            nn.ELU()
        )

        self.net = nn.Sequential(
            nn.LazyLinear(out_features=512),
            nn.ELU(),
            nn.Linear(512, 256),
            nn.ELU(),
            nn.Linear(256, 128),
            nn.ELU(),
            nn.Linear(128, self.num_actions)
        )

        self.log_std_parameter = nn.Parameter(torch.zeros(self.num_actions))

    def compute(self, inputs, role):
        # camera (N,H,W,C) -> (N,C,H,W)

        states = unflatten_tensorized_space(self.observation_space, inputs["states"])


        camera = states["camera"].permute(0, 3, 1, 2)
        camera_features = self.features_extractor(camera)

        state_features = self.net_state(states["robot-state"])

        fused = torch.cat([camera_features, state_features], dim=-1)
        out = self.net(fused)

        return torch.tanh(out), self.log_std_parameter, {}


class Value(DeterministicMixin, Model):
    def __init__(self, observation_space, action_space, device, clip_actions=True):
        Model.__init__(self, observation_space, action_space, device)
        DeterministicMixin.__init__(self, clip_actions)

        self.features_extractor = nn.Sequential(
            nn.LazyConv2d(out_channels=32, kernel_size=8, stride=4, padding=0),
            nn.ReLU(),
            nn.LazyConv2d(out_channels=64, kernel_size=4, stride=2, padding=0),
            nn.ReLU(),
            nn.LazyConv2d(out_channels=64, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Flatten()
        )
        
        self.net_state = nn.Sequential(
            nn.LazyLinear(out_features=128),
            nn.ELU(),
            nn.LazyLinear(out_features=128),
            nn.ELU(),
            nn.LazyLinear(out_features=64),
            nn.ELU()
        )

        self.net = nn.Sequential(
            nn.LazyLinear(out_features=512),
            nn.ELU(),
            nn.Linear(512, 256),
            nn.ELU(),
            nn.Linear(256, 128),
            nn.ELU(),
            nn.Linear(128, 1)
        )

    def compute(self, inputs, role):
        states = unflatten_tensorized_space(self.observation_space, inputs["states"])


        camera = states["camera"].permute(0, 3, 1, 2)
        camera_features = self.features_extractor(camera)

        state_features = self.net_state(states["robot-state"])

        fused = torch.cat([camera_features, state_features], dim=-1)
        return self.net(fused), {}

    
import gymnasium as gym

gym.register(
    id="Template-Vae-Navigation-Direct-v0",
    entry_point="VAE_navigation.tasks.direct.vae_navigation.vae_navigation_env:VaeNavigationEnv",
    disable_env_checker=True,
    kwargs={
        "env_cfg_entry_point": "VAE_navigation.tasks.direct.vae_navigation.vae_navigation_env_cfg:VaeNavigationEnvCfg",
        # "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml",
    },
)
env = load_isaaclab_env(task_name="Template-Vae-Navigation-Direct-v0", num_envs=1, headless=False)
env = wrap_env(env)

device = env.device

# instantiate a memory as rollout buffer (any memory can be used for this)
memory = RandomMemory(memory_size=512, num_envs=env.num_envs, device=device)

# instantiate the agent's models (function approximators).
models = {}
models["policy"] = Policy(
    observation_space=env.observation_space,
    action_space=env.action_space,
    device=device,
    # num_envs=env.num_envs,
)
models["value"] = Value(
    observation_space=env.observation_space,
    action_space=env.action_space,
    device=device,
    # num_envs=env.num_envs,
)

# configure and instantiate the agent
cfg = PPO_DEFAULT_CONFIG.copy()
cfg["rollouts"] = 512  # memory_size
cfg["learning_epochs"] = 8
cfg["mini_batches"] = 8
cfg["discount_factor"] = 0.99
cfg["lambda"] = 0.95
cfg["learning_rate"] = 1.0e-04
cfg["learning_rate_scheduler"] = KLAdaptiveRL
cfg["learning_rate_scheduler_kwargs"] = {"kl_threshold": 0.008}
cfg["random_timesteps"] = 0
cfg["learning_starts"] = 0
cfg["grad_norm_clip"] = 1.0
cfg["ratio_clip"] = 0.2
cfg["value_clip"] = 0.2
cfg["clip_predicted_values"] = True
cfg["entropy_loss_scale"] = 0.0
cfg["value_loss_scale"] = 2.0
cfg["kl_threshold"] = 0
cfg["rewards_shaper"] = lambda rewards, *args, **kwargs: rewards * 0.1
cfg["time_limit_bootstrap"] = True
cfg["state_preprocessor"] = RunningStandardScaler
cfg["state_preprocessor_kwargs"] = {"size": env.observation_space, "device": device}
cfg["value_preprocessor"] = RunningStandardScaler
cfg["value_preprocessor_kwargs"] = {"size": 1, "device": device}
cfg["experiment"]["write_interval"] = 1
cfg["experiment"]["checkpoint_interval"] = 2500
cfg["experiment"]["directory"] = "runs/torch/VAE"
cfg["experiment"]["wandb"] = True


agent = PPO(
    models=models,
    memory=memory,
    cfg=cfg,
    observation_space=env.observation_space,
    action_space=env.action_space,
    device=device,
)
# configure and instantiate the RL trainer
cfg_trainer = {"timesteps": 100000, "headless": False}
trainer = SequentialTrainer(cfg=cfg_trainer, env=env, agents=agent)

# start training
trainer.train()
# simulation_app.close()

and the error

  File "/home/trash/miniconda3/envs/env_isaaclab/lib/python3.10/site-packages/skrl/agents/torch/ppo/ppo.py", line 299, in record_transition
    rewards += self._discount_factor * values * truncated
RuntimeError: output with shape [1, 1] doesn't match the broadcast shape [1, 2]
[skrl:INFO] Closing environment

Can you help me with this bug ?
Thank you !

[Toni-SM]

Hi @giangdao1402

The Value model output should not be clipped, since it does not generate actions but an estimate of the value function.
Furthermore, setting clip_actions=True generate clipping values using the actions space size (self.num_actions), which may be different from the Value's final linear layer feature: 1.