jaxdem.rl.action_spaces.blockwise_space

Apply action-space bijectors to contiguous parts of one action vector.

Classes

BlockwiseSpace(*args, **kwargs)

Apply separate constraints to explicit slices of a flat action vector.

class jaxdem.rl.action_spaces.blockwise_space.BlockwiseSpace(*args, **kwargs)

Bases: Bijector, ActionSpace

Apply separate constraints to explicit slices of a flat action vector.

Parameters:

• blocks (sequence of (int, ActionSpace)) – Positive block sizes and their scalar or vector bijectors, in order. Sizes must sum to the input vector’s length. Each child must preserve its event shape and implement ActionSpace.log_det_expectation to support Gaussian policy entropy. Blocks may themselves be BlockwiseSpace instances.

• diagonal (The Jacobian is block)

• the (so its log-determinant is the sum of)

• likewise (children's log-determinants. The Gaussian entropy correction is)

• Gaussians. (the sum of their expectations over the corresponding marginal)

• preserved. (Leading sample and batch axes are)

Example

Constrain force and torque separately, followed by two bounded scalars:

BlockwiseSpace(blocks=(
    (3, MaxNormSpace(max_norm=force_limit)),
    (3, MaxNormSpace(max_norm=torque_limit)),
    (2, BoxSpace(-1.0, 1.0)),
))

property blocks: tuple[tuple[int, ActionSpace], ...]

Block sizes and original children, retained for factory serialization.

property size: int

Number of coordinates in one input or output event.

forward_and_log_det(x: Array | ndarray | bool | number) → tuple[Array, Array]

Computes y = f(x) and log|det J(f)(x)|.

inverse_and_log_det(y: Array | ndarray | bool | number) → tuple[Array, Array]

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

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

Sum the children’s Gaussian log-Jacobian expectations.