jaxdem.rl.action_spaces.free_space#

Identity bijector that applies no constraint.

Classes

FreeSpace(*args, **kwargs)

Identity constraint (no transform).
class jaxdem.rl.action_spaces.free_space.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.