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 entropy support for action-space bijectors.

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]

log_det_expectation(mean: Array, std: Array) → Array

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.Transformed(*args, **kwargs)

Bases: Transformed

distrax.Transformed with entropy support for action-space bijectors.

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

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

This is the differential-entropy identity for an invertible differentiable transform. For the diagonal-Gaussian policies used here, the bijector’s log_det_expectation() supplies the correction. FreeSpace and MaxNormSpace return it analytically. BoxSpace uses componentwise Gauss–Hermite quadrature, whose accuracy depends on the policy parameters and quadrature order. BlockwiseSpace sums its children’s corrections.

sample_and_log_prob_with_latent(*, seed: Array, sample_shape: tuple[int, ...] = ()) → tuple[Array, Array, Array, Array]

Sample an action and retain its latent and base log-probability.

PPO can compute same-transform likelihood ratios in base coordinates, where the Jacobian cancels, without inverting a rounded/saturated action. Returns action, transformed log-probability, latent, base log-probability.

entropy(input_hint: Array | ndarray | bool | number | None = None) → Array

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

blockwise_space	Apply action-space bijectors to contiguous parts of one action vector.
box_space	Bijector that constrains actions elementwise to a box.
free_space	Identity bijector that applies no constraint.
max_norm_space	Bijector that maps Gaussian actions into a ball with analytic entropy.