Add PPO critic lifecycle hooks and DistributionalPPO

rsl_rl/algorithms/__init__.py

@@ -6,6 +6,7 @@
 """Learning algorithms."""
 
 from .distillation import Distillation
+from .distributional_ppo import DistributionalPPO
 from .ppo import PPO
 
-__all__ = ["PPO", "Distillation"]
+__all__ = ["PPO", "Distillation", "DistributionalPPO"]

rsl_rl/algorithms/distributional_ppo.py

@@ -0,0 +1,148 @@
+# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+
+from __future__ import annotations
+
+import torch
+from tensordict import TensorDict
+
+from rsl_rl.env import VecEnv
+from rsl_rl.models import MLPModel
+from rsl_rl.storage import RolloutStorage
+
+from .ppo import PPO
+
+
+def quantile_huber_loss(quantiles: torch.Tensor, targets: torch.Tensor, kappa: float = 1.0) -> torch.Tensor:
+    """Compute Quantile Huber loss between predicted quantiles and scalar targets.
+
+    Reference:
+        - Dabney et al. "Distributional Reinforcement Learning with Quantile Regression." AAAI (2018).
+
+    Args:
+        quantiles: Predicted quantiles tensor of shape (..., num_quantiles).
+        targets: Target values tensor of shape (..., 1) or (...).
+        kappa: Huber loss threshold parameter (default: 1.0).
+
+    Returns:
+        Scalar quantile Huber loss tensor.
+    """
+    if targets.dim() == 1 or targets.shape[-1] != 1:
+        targets = targets.unsqueeze(-1)
+
+    num_quantiles = quantiles.shape[-1]
+    tau = (torch.arange(num_quantiles, device=quantiles.device, dtype=quantiles.dtype) + 0.5) / num_quantiles
+    tau = tau.unsqueeze(0)
+
+    pairwise_errors = targets - quantiles
+    abs_errors = torch.abs(pairwise_errors)
+
+    huber_loss = torch.where(
+        abs_errors <= kappa,
+        0.5 * pairwise_errors**2,
+        kappa * (abs_errors - 0.5 * kappa),
+    )
+
+    check_weight = torch.abs(tau - (pairwise_errors < 0).to(quantiles.dtype))
+    loss = check_weight * huber_loss
+
+    return torch.mean(loss)
+
+
+class DistributionalPPO(PPO):
+    """Distributional PPO using Quantile Regression and Huber Loss.
+
+    Extends standard PPO by estimating return distributions via quantiles
+    rather than a single scalar expected value. Scalar value estimates for GAE
+    are obtained by reducing over quantiles.
+
+    Reference:
+        - Dabney et al. "Distributional Reinforcement Learning with Quantile Regression." AAAI (2018).
+    """
+
+    def __init__(
+        self,
+        actor: MLPModel,
+        critic: MLPModel,
+        storage: RolloutStorage,
+        num_quantiles: int = 8,
+        quantile_huber_kappa: float = 1.0,
+        num_learning_epochs: int = 5,
+        num_mini_batches: int = 4,
+        clip_param: float = 0.2,
+        gamma: float = 0.998,
+        lam: float = 0.95,
+        value_loss_coef: float = 1.0,
+        entropy_coef: float = 0.01,
+        learning_rate: float = 0.001,
+        max_grad_norm: float = 1.0,
+        optimizer: str = "adam",
+        use_clipped_value_loss: bool = True,
+        schedule: str = "adaptive",
+        desired_kl: float = 0.01,
+        normalize_advantage_per_mini_batch: bool = False,
+        use_mixed_precision: bool = False,
+        device: str = "cpu",
+        # RND parameters
+        rnd_cfg: dict | None = None,
+        # Symmetry parameters
+        symmetry_cfg: dict | None = None,
+        # Distributed training parameters
+        multi_gpu_cfg: dict | None = None,
+        grad_reduce_bucket_mb: float = 25,
+    ) -> None:
+        """Initialize DistributionalPPO with quantile regression settings."""
+        super().__init__(
+            actor=actor,
+            critic=critic,
+            storage=storage,
+            num_learning_epochs=num_learning_epochs,
+            num_mini_batches=num_mini_batches,
+            clip_param=clip_param,
+            gamma=gamma,
+            lam=lam,
+            value_loss_coef=value_loss_coef,
+            entropy_coef=entropy_coef,
+            learning_rate=learning_rate,
+            max_grad_norm=max_grad_norm,
+            optimizer=optimizer,
+            use_clipped_value_loss=use_clipped_value_loss,
+            schedule=schedule,
+            desired_kl=desired_kl,
+            normalize_advantage_per_mini_batch=normalize_advantage_per_mini_batch,
+            use_mixed_precision=use_mixed_precision,
+            device=device,
+            rnd_cfg=rnd_cfg,
+            symmetry_cfg=symmetry_cfg,
+            multi_gpu_cfg=multi_gpu_cfg,
+            grad_reduce_bucket_mb=grad_reduce_bucket_mb,
+        )
+        self.num_quantiles = num_quantiles
+        self.quantile_huber_kappa = quantile_huber_kappa
+
+    @staticmethod
+    def construct_algorithm(obs: TensorDict, env: VecEnv, cfg: dict, device: str) -> DistributionalPPO:
+        """Construct the DistributionalPPO algorithm with quantile critic output dimension."""
+        cfg["critic"].setdefault("output_dim", cfg["algorithm"].get("num_quantiles", 8))
+        return PPO.construct_algorithm(obs, env, cfg, device)  # type: ignore[return-value]
+
+    def compute_value_estimate(self, obs: TensorDict) -> torch.Tensor:
+        """Compute mean scalar value estimate from quantiles for GAE advantage estimation."""
+        critic_output = super().compute_value_estimate(obs)
+        if self.num_quantiles > 0:
+            return critic_output.mean(dim=-1, keepdim=True)
+        return critic_output
+
+    def compute_critic_loss(
+        self,
+        value_preds: torch.Tensor,
+        returns: torch.Tensor,
+        old_value_preds: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute Quantile Huber loss over predicted quantiles."""
+        if self.num_quantiles > 0:
+            return quantile_huber_loss(value_preds, returns, kappa=self.quantile_huber_kappa)
+        return super().compute_critic_loss(value_preds, returns, old_value_preds)

rsl_rl/algorithms/ppo.py

@@ -127,7 +127,7 @@ def act(self, obs: TensorDict) -> torch.Tensor:
         self.transition.hidden_states = (self.actor.get_hidden_state(), self.critic.get_hidden_state())
         # Compute the actions and values
         self.transition.actions = self.actor(obs, stochastic_output=True).detach()
-        self.transition.values = self.critic(obs).detach()
+        self.transition.values = self.compute_value_estimate(obs)
         self.transition.actions_log_prob = self.actor.get_output_log_prob(self.transition.actions).detach()  # type: ignore
         self.transition.distribution_params = tuple(p.detach() for p in self.actor.output_distribution_params)
         # Record observations before env.step()
@@ -168,7 +168,7 @@ def compute_returns(self, obs: TensorDict) -> None:
         st = self.storage
         # Compute values for the last step
         critic_hidden_state = self.critic.get_hidden_state()
-        last_values = self.critic(obs).detach()
+        last_values = self.compute_value_estimate(obs)
         # Restore the critic's hidden state so the next rollout is not affected by the forward pass
         self.critic.reset(hidden_state=critic_hidden_state)
         # Compute returns and advantages
@@ -190,6 +190,48 @@ def compute_returns(self, obs: TensorDict) -> None:
         if not self.normalize_advantage_per_mini_batch:
             st.advantages = (st.advantages - st.advantages.mean()) / (st.advantages.std() + 1e-8)
 
+    def compute_value_estimate(self, obs: TensorDict) -> torch.Tensor:
+        """Compute scalar value estimate from observations.
+
+        Subclasses (e.g. distributional critics or ensemble critics) can override
+        this method to customize how scalar value estimates are computed for
+        GAE advantage estimation.
+
+        Args:
+            obs: Observation dictionary.
+
+        Returns:
+            Scalar value estimate tensor.
+        """
+        return self.critic(obs).detach()
+
+    def compute_critic_loss(
+        self,
+        value_preds: torch.Tensor,
+        returns: torch.Tensor,
+        old_value_preds: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute value function loss.
+
+        Subclasses can override this method to implement custom critic objectives
+        (e.g., Quantile Huber loss for distributional critics).
+
+        Args:
+            value_preds: Current value predictions from the critic.
+            returns: Target returns.
+            old_value_preds: Value predictions from rollout before update.
+
+        Returns:
+            Scalar value loss tensor.
+        """
+        if self.use_clipped_value_loss:
+            value_clipped = old_value_preds + (value_preds - old_value_preds).clamp(-self.clip_param, self.clip_param)
+            value_losses = (value_preds - returns).pow(2)
+            value_losses_clipped = (value_clipped - returns).pow(2)
+            return torch.max(value_losses, value_losses_clipped).mean()
+        else:
+            return (returns - value_preds).pow(2).mean()
+
     def update(self) -> dict[str, float]:
         """Run optimization epochs over stored batches and return mean losses."""
         mean_value_loss = 0
@@ -274,13 +316,7 @@ def update(self) -> dict[str, float]:
                 surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()
 
                 # Value function loss
-                if self.use_clipped_value_loss:
-                    value_clipped = batch.values + (values - batch.values).clamp(-self.clip_param, self.clip_param)
-                    value_losses = (values - batch.returns).pow(2)
-                    value_losses_clipped = (value_clipped - batch.returns).pow(2)
-                    value_loss = torch.max(value_losses, value_losses_clipped).mean()
-                else:
-                    value_loss = (batch.returns - values).pow(2).mean()
+                value_loss = self.compute_critic_loss(values, batch.returns, batch.values)
 
                 loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy.mean()
 
@@ -448,7 +484,9 @@ def construct_algorithm(obs: TensorDict, env: VecEnv, cfg: dict, device: str) ->
         print(f"Actor Model: {actor}")
         if alg_cfg.pop("share_cnn_encoders", None):  # Share CNN encoders between actor and critic
             critic_cfg["cnns"] = actor.cnns
-        critic: MLPModel = critic_class(obs, cfg["obs_groups"], "critic", 1, **critic_cfg).to(device)
+        critic: MLPModel = critic_class(
+            obs, cfg["obs_groups"], "critic", critic_cfg.pop("output_dim", 1), **critic_cfg
+        ).to(device)
         print(f"Critic Model: {critic}")
 
         # Initialize the storage