jaxdem.rl.actionSpaces.freeSpace#

Implementation of identity bijector for free space.

Classes

FreeSpace(*args, **kwargs)

Identity constraint (no transform).
class jaxdem.rl.actionSpaces.freeSpace.FreeSpace(*args, **kwargs)[source][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 (defaults to 0 for a scalar transform).

  • -event_ndims_out (Optional[int]) – standard Distrax/TFP bijector flags.

  • -is_constant_jacobian (bool) – standard Distrax/TFP bijector flags.

  • -is_constant_log_det (bool) – standard Distrax/TFP bijector flags.

Note

This bijector is scalar (event_ndims_in = 0). For vector actions, needs to be wrap it with ``distrax.Block(bijector, ndims=1)`. Let the model do that for you!

property kws: Dict[source]#
forward_and_log_det(x: Array | ndarray | bool | number) → Tuple[Array | ndarray | bool | number, Array][source][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][source]#

Computes x = f^{-1}(y) and log|det J(f^{-1})(y)|.

same_as(other: Bijector) → bool[source][source]#

Returns True if this bijector is guaranteed to be the same as other.

classmethod registry_name() → str[source]#
property type_name: str[source]#