jaxdem.rl.action_spaces.box_space#

Bijector that constrains actions elementwise to a box.

Classes

BoxSpace(*args, **kwargs)

Elementwise box constraint implemented with a scaled tanh.

class jaxdem.rl.action_spaces.box_space.BoxSpace(*args, **kwargs)[source]#

Bases: Bijector, ActionSpace

Elementwise box constraint implemented with a scaled tanh.

Mapping (componentwise)

\[y_i \;=\; c_i + h_i\,\tanh\!\left(\frac{x_i}{w}\right), \qquad c_i=\tfrac{1}{2}(x_{\min,i}+x_{\max,i}), \quad h_i=\tfrac{1-\varepsilon}{2}(x_{\max,i}-x_{\min,i}),\]

with a width parameter \(w>0\) and a small \(\epsilon>0\) for numerical safety.

Jacobian (componentwise) For each component,

\[\frac{\partial y_i}{\partial x_i} = \frac{h_i}{w} sech^2 \left(\frac{x_i}{w}\right), \qquad \log\left| \frac{\partial y_i}{\partial x_i} \right| = \log h_i - \log w + \log\!\big(sech^2(\frac{x_i}{w})\big).\]

We use the stable identity \(\log(sech^2 z)=2 [\log 2 - z - softplus(-2z)]\) for good numerical behavior.

Parameters:

x_min (jax.Array) – Elementwise lower bounds of the box.
x_max (jax.Array) – Elementwise upper bounds of the box. Must satisfy x_max > x_min elementwise.
width (float) – Controls the tanh slope (default 1.0).
eps (float) – Small offset to avoid arctanh divergence close to the bounds (default 1e-6).
event_ndims_in (int) – Dimensionality of a single event seen by the bijector (default 0 for a scalar transform).
event_ndims_out (Optional[int]) – Standard Distrax/TFP bijector flag.
is_constant_jacobian (bool) – Standard Distrax/TFP bijector flag.
is_constant_log_det (bool) – Standard Distrax/TFP bijector flag.

Note

This bijector is scalar (event_ndims_in = 0). For vector actions, wrap it with distrax.Block(bijector, ndims=1). The model applies this wrapper automatically.

static sec2_log(x: Array) → Array[source]#

forward_log_det_jacobian(x: Array | ndarray | bool | number) → Array[source]#: Compute log|det J(f)(x)| = log(half) - log(width) + log(sech^2(x/width)). Uses the stable identity log(sech^2 z) = 2*(log(2) - z - softplus(-2z)).

forward_and_log_det(x: Array | ndarray | bool | number) → tuple[Array, Array][source]#: Compute y = f(x) and log|det J(f)(x)|.

inverse_and_log_det(y: Array | ndarray | bool | number) → tuple[Array, Array][source]#: Compute x = f^{-1}(y) and log|det J(f^{-1})(y)|.

same_as(other: Bijector) → bool[source]#: Return True if this bijector is guaranteed to be the same as other.

log_det_expectation(mean: Array, std: Array) → Array[source]#: \(\mathbb{E}_X[\sum_i \log|dJ_i/dx_i|]\) via 1-D Gauss-Hermite quadrature (componentwise separable).