jaxdem.rl.action_spaces#

Interface for bijectors that constrain the policy probability distribution.

Classes

`ActionSpace`()	Registry/namespace for action-space constraints implemented as `distrax.Bijector` objects.
`Transformed`(args, *kwargs)	`distrax.Transformed` with analytical entropy support.

class jaxdem.rl.action_spaces.ActionSpace#

Bases: Factory

Registry/namespace for action-space constraints implemented as distrax.Bijector objects.

Wrap these bijectors around a base policy distribution (e.g., MultivariateNormalDiag) with distrax.Transformed. The bijector’s ‘forward_and_log_det’ / ‘inverse_and_log_det’ methods then adjust sampling and log-probabilities correctly. See the Distrax/TFP bijector interface for details on shape semantics and ‘event_ndims_in/out’.

Example:#

To define a custom action space, inherit from distrax.Bijector and ActionSpace and implement its abstract methods:

>>> @ActionSpace.register("myCustomActionSpace")
>>> class MyCustomActionSpace(distrax.Bijector, ActionSpace):
        ...

property kws: dict[str, Any][source]#

log_det_expectation(mean: Array, std: Array) → Array[source]#

Compute \(\mathbb{E}_{X}[\log|\det J_f(X)|]\) where \(X \sim \mathcal{N}(\text{mean}, \text{diag}(\text{std}^2))\).

Subclasses override this method to enable Transformed.entropy().

class jaxdem.rl.action_spaces.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).

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

Bases: Bijector, ActionSpace

Identity constraint (no transform).

Mapping

\[y = f(x) = x, \qquad x = f^{-1}(y) = y.\]

Jacobian

\[J_f(x) = I,\qquad \log\lvert\det J_f(x)\rvert = 0, \qquad \log\lvert\det J_{f^{-1}}(y)\rvert = 0.\]

Parameters:

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.

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

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

log_det_expectation(mean: Array, std: Array) → Array[source]#

Compute \(\mathbb{E}_{X}[\log|\det J_f(X)|]\) where \(X \sim \mathcal{N}(\text{mean}, \text{diag}(\text{std}^2))\).

Subclasses override this method to enable Transformed.entropy().

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

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

Bases: Bijector, ActionSpace

Radial max-norm constraint for vector actions. Scales the radius with a tanh squashing and preserves the direction.

Mapping (vector case), \(\vec{x} \in \mathbb{R}^d\):

\[\begin{split}r = \lVert \vec{x} \rVert_2,\qquad \hat{u} = \begin{cases} \frac{\vec{x}}{r}, & r>0,\\[4pt] 0, & r=0, \end{cases} \qquad y = s \tanh(r) \hat{u}, \quad s = (1-\epsilon) \text{max\_norm}.\end{split}\]

Equivalently, \(y = b(r)\,x\) with \(b(r)= s\,\tanh(r)/r\) for \(r>0\).

Jacobian determinant

For an isotropic radial map \(f(x)=b(r)\) with \(x \in \mathbb{R}^d\), the Jacobian eigenvalues are \(b\) (multiplicity d-1) on the tangent subspace and \(b + r\,b'(r)\) on the radial direction. Therefore

\[\bigl|\det J_f(x)\bigr| = b(r)^{\,d-1}\,\bigl(b(r)+r\,b'(r)\bigr) = s^d \left(\frac{\tanh r}{r}\right)^{\!d-1} \text{sech}^2 r.\]

Therefore

\[\log\lvert\det J_f(x)\rvert = d\log s + (d-1)\bigl(\log\tanh r - \log r\bigr) + \log(\text{sech}^2 r),\]

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

Near \(r\approx 0\), we use the second-order expansion

\[\log\lvert\det J_f(x)\rvert \approx d\log s - \tfrac{2}{3} r^2\]

to avoid division by \(r\).

Parameters:

max_norm (float) – Maximum radius after squashing (default 1.0). The bijector uses (s=(1-varepsilon),text{max_norm}) to stay off the exact boundary.
eps (float) – Numerical safety margin near (r=0) and (rtoinfty).
event_ndims_in (int) – Dimensionality of a single event seen by the bijector (default 1).
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 vector-valued with event_ndims_in = 1. It treats a length-(d) action vector as a single event. Do not wrap it in Block unless you want to apply it independently to multiple last-axis blocks.

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

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)|.

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

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

log_det_expectation(mean: Array, std: Array) → Array[source]#: \(\mathbb{E}_X[\log|\det J_f(X)|]\) via tensor-product Gauss-Hermite quadrature in d dimensions.

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

Bases: Transformed

`distrax.Transformed` with analytical entropy support.

For \(Y = f(X)\) where \(X \sim \text{base}\),

\[H(Y) = H(X) + \mathbb{E}_X[\log|\det J_f(X)|].\]

The bijector’s log_det_expectation() method computes the expectation with Gauss–Hermite quadrature. The quadrature is exact for polynomial integrands and accurate for smooth bijectors such as scaled tanh.

entropy(input_hint: Array | ndarray | bool | number | None = None) → Array[source]#

Calculates the Shannon entropy (in Nats).

Only works for bijectors with constant Jacobian determinant.

Parameters:: input_hint – an example sample from the base distribution, used to compute the constant forward log-determinant. If not specified, it is computed using a zero array of the shape and dtype of a sample from the base distribution.
Returns:: the entropy of the distribution.
Raises:: NotImplementedError – if bijector’s Jacobian determinant is not known to be constant.

Modules

`box_space`	Bijector that constrains actions elementwise to a box.
`free_space`	Identity bijector that applies no constraint.
`max_norm_space`	Bijector that constrains the norm of vector actions.