# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
"""Value normalisation for actor-critic algorithms.
Defines the abstract :class:`ValueNorm` interface and two concrete
implementations:
- :class:`PopArtValueNorm` — exponential-moving-average mean / mean-of-squares
with debiasing (van Hasselt et al., *Multi-task Deep RL with PopArt*,
AAAI 2019, https://arxiv.org/abs/1809.04474). Used by MAPPO
(Yu et al. 2022) to stabilise the critic loss when reward scales drift.
- :class:`RunningValueNorm` — exact Welford running mean / variance with no
decay. Cheaper and more stable when value targets are stationary; tends to
be the better default for shorter / non-curriculum runs.
- :class:`PercentileValueNorm` — exponential-moving-average of a low / high
quantile pair, rescaling by the clamped span between them (the DreamerV3
"return normalisation", Hafner et al. 2023,
https://arxiv.org/abs/2301.04104). Scale-only by default, which suits
advantage scaling.
Plug any subclass into :class:`~torchrl.objectives.multiagent.MAPPOLoss` (or
your own actor-critic loss) via ``value_norm=...``.
"""
from __future__ import annotations
import math
import warnings
from abc import ABCMeta, abstractmethod
import torch
from torch import nn
[docs]
class ValueNorm(nn.Module, metaclass=ABCMeta):
"""Abstract base class for value normalisers.
A *value normaliser* keeps a running estimate of the location and scale of
the value target seen during training. Critics use it to:
- **normalize** the regression target before computing MSE, keeping the
critic loss on a fixed scale across episodes / reward inflations;
- **denormalize** the critic's output back to the real reward scale when
forming bootstrapped value estimates inside GAE / TD.
Subclasses must implement :meth:`update`, :meth:`normalize`,
:meth:`denormalize`, and :meth:`scale`. The convention is that they all
operate on tensors whose trailing dims match :attr:`shape` (the
per-element value shape, usually ``(1,)``).
"""
shape: tuple[int, ...]
def __init__(
self,
*,
shape: int | tuple[int, ...] = 1,
epsilon: float = 1e-5,
device: torch.device | None = None,
) -> None:
super().__init__()
if isinstance(shape, int):
shape = (shape,)
self.shape = tuple(shape)
self.epsilon = epsilon
self._device = device
if type(self).scale is ValueNorm.scale:
warnings.warn(
f"{type(self).__name__} does not implement scale(). "
"ValueNorm.scale() will become abstract in v0.16; implement it "
"in your subclass to keep it instantiable.",
DeprecationWarning,
stacklevel=2,
)
# ------------------------------------------------------------------ API
[docs]
@abstractmethod
def update(self, value_target: torch.Tensor) -> None:
"""Fold a batch of value targets into the running stats."""
[docs]
@abstractmethod
def normalize(self, value_target: torch.Tensor) -> torch.Tensor:
"""Standardise ``value_target`` using the current running stats."""
[docs]
@abstractmethod
def denormalize(self, normalised_value: torch.Tensor) -> torch.Tensor:
"""Inverse of :meth:`normalize` — recover real-scale values."""
[docs]
def scale(self) -> torch.Tensor:
"""Multiplicative scale currently applied by :meth:`normalize`.
Exposed separately so consumers can rescale quantities that must not
be re-centred, e.g. advantages (already centred by the value
baseline), for which only the division by the scale applies.
Deliberately not abstract until v0.16 so that subclasses written
before it existed keep instantiating (with a ``DeprecationWarning``);
this default raises ``NotImplementedError`` when called.
"""
raise NotImplementedError(
f"{type(self).__name__} does not implement scale(). Implement it "
"to use this normaliser where a multiplicative scale is required "
"(e.g. advantage scaling)."
)
# ------------------------------------------------------- shared helpers
def _check_trailing_shape(self, value_target: torch.Tensor) -> tuple[int, ...]:
if value_target.shape[-len(self.shape) :] != self.shape:
raise ValueError(
f"{type(self).__name__} was initialised with shape={self.shape} "
f"but got a value_target with trailing shape "
f"{tuple(value_target.shape[-len(self.shape) :])}."
)
return tuple(range(value_target.ndim - len(self.shape)))
[docs]
class PopArtValueNorm(ValueNorm):
"""PopArt-style EMA value normaliser.
Maintains exponentially-weighted running estimates of the value-target
mean and mean-of-squares, with debiasing (so the early-training estimates
are unbiased even before the EMA has had time to wash out the zero
initialisation). Equivalent to the value-normaliser used by the reference
MAPPO implementation.
Keyword Args:
shape: per-element shape of the value tensor (everything except the
leading batch / time / agent dims that get reduced). Defaults to
``1``.
beta: exponential decay for the running stats. Higher = slower
adaptation. Defaults to ``0.99999`` (the MAPPO default).
epsilon: numerical stabiliser added to the running variance and used
as a floor for the debiasing term. Defaults to ``1e-5``.
device: device for the running-stats buffers.
Example:
>>> vn = PopArtValueNorm(shape=1)
>>> target = torch.randn(64, 1) * 5.0 + 2.0 # mean 2, std 5
>>> for _ in range(100):
... vn.update(target)
>>> normed = vn.normalize(target) # ~ N(0, 1)
>>> recovered = vn.denormalize(normed) # back to real scale
"""
def __init__(
self,
*,
shape: int | tuple[int, ...] = 1,
beta: float = 0.99999,
epsilon: float = 1e-5,
device: torch.device | None = None,
) -> None:
super().__init__(shape=shape, epsilon=epsilon, device=device)
self.beta = beta
# Both running buffers start at zero. The debiasing term tracks
# \sum_{s<=t} beta^{t-s}, which also starts at zero; dividing the
# zero-init buffers by the (clamped) debias gives an unbiased EMA.
self.register_buffer("running_mean", torch.zeros(self.shape, device=device))
self.register_buffer("running_mean_sq", torch.zeros(self.shape, device=device))
self.register_buffer("debiasing_term", torch.zeros((), device=device))
def _running_stats(self) -> tuple[torch.Tensor, torch.Tensor]:
debias = self.debiasing_term.clamp(min=self.epsilon)
mean = self.running_mean / debias
mean_sq = self.running_mean_sq / debias
var = (mean_sq - mean.pow(2)).clamp(min=self.epsilon)
return mean, var
[docs]
@torch.no_grad()
def update(self, value_target: torch.Tensor) -> None:
value_target = value_target.detach()
reduce_dims = self._check_trailing_shape(value_target)
if reduce_dims:
batch_mean = value_target.mean(dim=reduce_dims)
batch_mean_sq = value_target.pow(2).mean(dim=reduce_dims)
else:
batch_mean = value_target
batch_mean_sq = value_target.pow(2)
self.running_mean.mul_(self.beta).add_(batch_mean, alpha=1.0 - self.beta)
self.running_mean_sq.mul_(self.beta).add_(batch_mean_sq, alpha=1.0 - self.beta)
self.debiasing_term.mul_(self.beta).add_(1.0 - self.beta)
[docs]
def normalize(self, value_target: torch.Tensor) -> torch.Tensor:
mean, var = self._running_stats()
return (value_target - mean) / var.sqrt()
[docs]
def denormalize(self, normalised_value: torch.Tensor) -> torch.Tensor:
mean, var = self._running_stats()
return normalised_value * var.sqrt() + mean
[docs]
def scale(self) -> torch.Tensor:
_, var = self._running_stats()
return var.sqrt()
[docs]
class RunningValueNorm(ValueNorm):
"""Exact running mean / variance (Welford's online algorithm).
Unlike :class:`PopArtValueNorm`, this normaliser does not decay older
samples — it accumulates the true sample mean and variance over every
target it has ever seen. Useful when value targets are roughly stationary
(no curriculum, no reward-shaping schedule), where the EMA's adaptivity
is unnecessary and the exact running stats give a slightly tighter
estimate.
Keyword Args:
shape: per-element shape of the value tensor. Defaults to ``1``.
epsilon: numerical stabiliser added to the running variance.
Defaults to ``1e-5``.
device: device for the running-stats buffers.
Example:
>>> vn = RunningValueNorm(shape=1)
>>> for _ in range(10):
... vn.update(torch.randn(64, 1) * 3.0 + 1.0)
>>> normed = vn.normalize(torch.randn(8, 1))
"""
def __init__(
self,
*,
shape: int | tuple[int, ...] = 1,
epsilon: float = 1e-5,
device: torch.device | None = None,
) -> None:
super().__init__(shape=shape, epsilon=epsilon, device=device)
self.register_buffer("mean", torch.zeros(self.shape, device=device))
# m2 stores the running sum of squared deviations from the mean
# (Welford's M2). var = m2 / count (population variance — we don't
# apply Bessel's correction because the consumer is normalisation,
# not statistical inference).
self.register_buffer("m2", torch.zeros(self.shape, device=device))
self.register_buffer("count", torch.zeros((), device=device))
[docs]
@torch.no_grad()
def update(self, value_target: torch.Tensor) -> None:
value_target = value_target.detach()
reduce_dims = self._check_trailing_shape(value_target)
if reduce_dims:
# ``math.prod`` over plain Python ints avoids allocating a
# CPU tensor and forcing a host-device sync inside update().
batch_count = float(math.prod(value_target.shape[d] for d in reduce_dims))
batch_mean = value_target.mean(dim=reduce_dims)
batch_var = value_target.var(dim=reduce_dims, unbiased=False)
else:
batch_count = 1.0
batch_mean = value_target
batch_var = torch.zeros_like(value_target)
# Chan et al. parallel variance update.
delta = batch_mean - self.mean
total = self.count + batch_count
new_mean = self.mean + delta * (batch_count / total)
new_m2 = (
self.m2
+ batch_var * batch_count
+ delta.pow(2) * (self.count * batch_count / total)
)
self.mean.copy_(new_mean)
self.m2.copy_(new_m2)
self.count.fill_(total)
def _var(self) -> torch.Tensor:
denom = self.count.clamp(min=1.0)
return (self.m2 / denom).clamp(min=self.epsilon)
[docs]
def normalize(self, value_target: torch.Tensor) -> torch.Tensor:
return (value_target - self.mean) / self._var().sqrt()
[docs]
def denormalize(self, normalised_value: torch.Tensor) -> torch.Tensor:
return normalised_value * self._var().sqrt() + self.mean
[docs]
def scale(self) -> torch.Tensor:
return self._var().sqrt()
[docs]
class PercentileValueNorm(ValueNorm):
"""DreamerV3-style EMA percentile-range value normaliser.
Tracks exponential moving averages of a low and a high quantile of the
value targets and rescales by the span between them, clamped from below:
``scale = max(min_scale, high - low)``. Following DreamerV3 (Hafner et
al., *Mastering Diverse Domains through World Models*, 2023,
https://arxiv.org/abs/2301.04104), the clamp scales large values down
without amplifying small or noisy ones, which keeps fixed coefficients
such as an entropy bonus comparable across reward scales.
By default (``center=False``) :meth:`normalize` only divides by the
span — the DreamerV3 recipe for advantages, which are already centred by
the value baseline. With ``center=True`` the low-percentile EMA is also
subtracted, mapping the tracked percentile range onto ``[0, 1]``.
Keyword Args:
shape: per-element shape of the value tensor (everything except the
leading batch / time / agent dims that get reduced). Defaults to
``1``.
quantiles: lower and upper quantiles tracked by the EMA. Defaults to
``(0.05, 0.95)``.
rate: EMA update rate towards the batch quantiles; higher = faster
adaptation. Defaults to ``0.01``.
min_scale: lower bound of the normalisation scale. Defaults to
``1.0``.
center: if ``True``, subtract the low-percentile EMA in
:meth:`normalize`. Defaults to ``False``.
epsilon: kept for interface parity with the other normalisers;
unused because ``min_scale`` already bounds the divisor.
device: device for the running-stats buffers.
Example:
>>> vn = PercentileValueNorm(shape=1, rate=1.0)
>>> returns = torch.linspace(0.0, 100.0, steps=101).unsqueeze(-1)
>>> vn.update(returns)
>>> vn.scale()
tensor([90.])
>>> vn.normalize(torch.tensor([45.0]))
tensor([0.5000])
"""
def __init__(
self,
*,
shape: int | tuple[int, ...] = 1,
quantiles: tuple[float, float] = (0.05, 0.95),
rate: float = 0.01,
min_scale: float = 1.0,
center: bool = False,
epsilon: float = 1e-5,
device: torch.device | None = None,
) -> None:
super().__init__(shape=shape, epsilon=epsilon, device=device)
low, high = quantiles
if not 0 <= low < high <= 1:
raise ValueError(
f"quantiles must satisfy 0 <= low < high <= 1, got {quantiles}."
)
if not 0 <= rate <= 1:
raise ValueError(f"rate must be in [0, 1], got {rate}.")
if min_scale <= 0:
raise ValueError(f"min_scale must be positive, got {min_scale}.")
self.quantiles = (float(low), float(high))
self.rate = rate
self.min_scale = min_scale
self.center = center
self.register_buffer("low", torch.zeros(self.shape, device=device))
self.register_buffer("high", torch.zeros(self.shape, device=device))
self.register_buffer(
"_q",
torch.tensor([float(low), float(high)], device=device),
persistent=False,
)
[docs]
@torch.no_grad()
def update(self, value_target: torch.Tensor) -> None:
value_target = value_target.detach()
self._check_trailing_shape(value_target)
flat = value_target.reshape(-1, *self.shape).to(self.low.dtype)
batch_low, batch_high = torch.quantile(
flat, self._q.to(device=flat.device, dtype=flat.dtype), dim=0
)
self.low.lerp_(batch_low, self.rate)
self.high.lerp_(batch_high, self.rate)
[docs]
def normalize(self, value_target: torch.Tensor) -> torch.Tensor:
if self.center:
return (value_target - self.low) / self.scale()
return value_target / self.scale()
[docs]
def denormalize(self, normalised_value: torch.Tensor) -> torch.Tensor:
if self.center:
return normalised_value * self.scale() + self.low
return normalised_value * self.scale()
[docs]
def scale(self) -> torch.Tensor:
return (self.high - self.low).clamp_min(self.min_scale)
