jaxdem.rl.action_spaces.box_space#
Bijector that constrains actions elementwise to a box.
Classes
|
Elementwise box constraint implemented with a scaled tanh. |
- class jaxdem.rl.action_spaces.box_space.BoxSpace(*args, **kwargs)[source]#
Bases:
Bijector,ActionSpaceElementwise 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) – Must be zero for this scalar transform.
event_ndims_out (Optional[int]) – Must be None or zero. Use
distrax.Blockfor vector events.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 withdistrax.Block(bijector, ndims=1). The model applies this wrapper automatically.- 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]#
Evaluate the inverse after clipping the normalized coordinate.
Only clips to the nearest representable points inside (-1, 1), independently of the output margin
eps. Saturated floating-point actions cannot recover their original latent exactly; PPO retains that latent during sampling instead of relying on this inverse.