jaxdem.rl.action_spaces.max_norm_space#
Bijector that maps Gaussian actions into a ball with analytic entropy.
Classes
|
Constrain the norm of one action vector using a Gaussian-to-ball map. |
- class jaxdem.rl.action_spaces.max_norm_space.MaxNormSpace(*args, **kwargs)[source]#
Bases:
Bijector,ActionSpaceConstrain the norm of one action vector using a Gaussian-to-ball map.
For an event of length d, preserve its direction and transform its radius:
\[r=\|x\|,\quad R=(1-\epsilon)\,\mathrm{max\_norm},\qquad y=R\,P(d/2,r^2/2)^{1/d}\frac{x}{r},\]where P is the regularized lower incomplete gamma function. This is the probability integral transform of a standard Gaussian’s chi-distributed radius, followed by the inverse radius CDF of a uniform ball. The reference Gaussian is fixed at zero mean and unit variance; learned Gaussian means and variances still control the action distribution.
Writing c for the linear scale at the origin gives
\[c=\frac{R}{\sqrt{2}\,\Gamma(d/2+1)^{1/d}},\qquad y\sim cx,\qquad \log|\det J(x)|=d\log c-\tfrac12\|x\|^2.\]Thus the Gaussian log-Jacobian expectation and transformed entropy are analytic, including for anisotropic, noncentral Gaussian policies:
\[\begin{split}\mathbb E[\log|\det J(X)|] =d\log c-\tfrac12\sum_i(\mu_i^2+\sigma_i^2),\\ H(Y)=\log V_d(R) -D_{\mathrm{KL}}(\mathcal N(\mu,\mathrm{diag}(\sigma^2)) \,\|\,\mathcal N(0,I)).\end{split}\]The forward and inverse use a convergent gamma series near zero to avoid dividing by the radius or forming an underflowing d-th power. Elsewhere, they use the regularized gamma function and its inverse. At the boundary, the inverse clips to the nearest representable interior squared radius. Saturated actions cannot recover their original latent exactly; PPO retains that latent when sampling.
- Parameters:
max_norm (float) – Positive finite upper bound on the output norm, default 1.0.
eps (float) – Output-radius margin in (0, 1), default 1e-6.
event_ndims_in (int) – Must be one for a single vector.
event_ndims_out (int or None) – Must be None or one. Use
distrax.Blockfor higher-rank events.is_constant_jacobian (bool) – Standard Distrax flag, default False.
is_constant_log_det (bool or None) – Standard Distrax flag, default None.
Notes
This replaces the historical radial tanh map. Existing Gaussian model heads and constructor arguments are supported, but the same latent now produces different actions. One instance constrains the entire last-axis vector; BlockwiseSpace applies separate bounds to explicit sub-vectors.
- forward_log_det_jacobian(x: Array | ndarray | bool | number) Array[source]#
Computes log|det J(f)(x)|.
- forward_and_log_det(x: Array | ndarray | bool | number) tuple[Array, Array][source]#
Computes y = f(x) and log|det J(f)(x)|.