jaxdem.utils.environment#

Utility functions to handle environments and LIDAR sensor.

Functions

advance_action(env, action, *[, ...])

Snapshot the action baseline once, then advance physics to its endpoint.

cross_lidar_2d(pos_a, pos_b, system, ...)

2-D LIDAR proximity and IDs from pos_a sensing targets in pos_b.

cross_lidar_3d(pos_a, pos_b, system, ...)

3-D LIDAR proximity and IDs from pos_a sensing targets in pos_b.

env_step(env, model, key, graphstate, *[, ...])

Advance the environment n steps using actions from model.

env_trajectory_rollout(env, model, key, ...)

Roll out a trajectory by applying model in chunks of stride steps and collecting the environment after each chunk.

lidar_2d(state, system, lidar_range, n_bins)

2-D LIDAR proximity readings and neighbor IDs.

lidar_3d(state, system, lidar_range, ...[, ...])

3-D LIDAR proximity readings and neighbor IDs.
jaxdem.utils.environment.advance_action(env: Environment, action: jax.Array, *, skip_frames: int = 0, terminated: jax.Array | None = None, truncated: jax.Array | None = None) → tuple[Environment, jax.Array, jax.Array][source]#

Snapshot the action baseline once, then advance physics to its endpoint.

Evaluates 1 + skip_frames physics steps for the full batch. Updates after each environment’s first termination/truncation are discarded, preserving its boundary state independently of the other batch entries. Calls env.checkpoint(env, action) before the first physics step and preserves that historical baseline throughout the action. Returns the live endpoint and the two boundary flags without a second checkpoint. It does not reset episodes or accumulate intermediate rewards. Read endpoint reward before starting another action or resetting the environment. By default this starts a new transition (the caller must reset completed episodes). Optional incoming boundary masks keep already finished members frozen when continuing an evaluation rollout without automatic resets. Inactive agent actions are zeroed before checkpointing and stepping, preserving integer categorical indices and floating-point continuous actions.

jaxdem.utils.environment.env_trajectory_rollout(env: Environment, model: Callable[..., Any], key: jax.Array, graphstate: Any, *, n: int, stride: int = 1, skip_frames: int = 0, **kw: Any) → tuple[Environment, jax.Array, Any, Environment][source]#

Roll out a trajectory by applying model in chunks of stride steps and collecting the environment after each chunk.

This performs n * stride policy decisions with the same action/key sequence as one env_step(..., n=n*stride) call. Recording boundaries do not consume additional random keys.

Parameters:

  • env (Environment) – Initial environment pytree.

  • model (Callable) – Callable with signature model(obs, key, graphstate, **kw) -> (action, graphstate).

  • key (jax.Array) – Initial random key.

  • graphstate (Any) – Initial model state.

  • n (int) – Total number of macro-steps.

  • stride (int) – Number of steps to perform before recording the environment state.

  • skip_frames (int) – Number of additional physics frames to repeat each action, so every logical step requests 1 + skip_frames physics frames. State updates after a boundary are discarded; completed episodes stay frozen until the caller resets them. Defaults to 0.

  • **kw (Any) – Extra keyword arguments passed to model on every step.

Returns:

Final environment, advanced random key, final graphstate, and a stacked pytree of environments with length n, each snapshot taken after a chunk of stride steps.

Return type:

Tuple[Environment, jax.Array, Any, Environment]

Examples

>>> env, key, graphstate, traj = env_trajectory_rollout(env, model, key, graphstate, n=100, stride=5, objective=goal)
jaxdem.utils.environment.env_step(env: Environment, model: Callable[..., Any], key: jax.Array, graphstate: Any, *, n: int = 1, skip_frames: int = 0, **kw: Any) → tuple[Environment, jax.Array, Any][source]#

Advance the environment n steps using actions from model.

Each policy action begins with one historical checkpoint. Observations and rewards read the resulting live state; no endpoint checkpoint is needed. Finished batch entries remain physically frozen, with their baseline refreshed for each requested action. This helper does not reset episodes or recurrent carry. The model is still called for each requested logical step, even for frozen environments; its state may therefore advance.

Parameters:

  • env (Environment) – Initial environment pytree (batchable).

  • model (Callable) – Callable with signature model(obs, key, graphstate, **kw) -> (action, graphstate).

  • key (jax.Array) – JAX random key. Use the returned (advanced) key for subsequent calls.

  • graphstate (Any) – Initial model state.

  • n (int) – Number of steps to perform.

  • skip_frames (int) – Number of additional physics frames to repeat each action, so every logical step requests 1 + skip_frames physics frames. State updates after a boundary are discarded; completed episodes stay frozen until the caller resets them. Defaults to 0.

  • **kw (Any) – Extra keyword arguments forwarded to model.

Returns:

Updated environment, the advanced random key, and updated graphstate.

Return type:

Tuple[Environment, jax.Array, Any]

Examples

>>> env, key, graphstate = env_step(env, model, key, graphstate, n=10, objective=goal)
jaxdem.utils.environment.lidar_2d(state: State, system: System, lidar_range: float, n_bins: int, sense_edges: bool = False) → tuple[State, System, jax.Array, jax.Array, jax.Array][source]#

2-D LIDAR proximity readings and neighbor IDs.

For every particle in state the function projects the displacement vectors to all other particles onto the \(xy\)-plane and bins them by azimuthal angle into n_bins uniform sectors spanning \([-\pi, \pi)\). Each bin stores the proximity value and the index of the closest neighbor in that sector:

\[p_k = \max(0,\; r_{\max} - d_{\min,k})\]

This works identically for 2-D and 3-D position data. In the 3-D case the function ignores the \(z\)-component of the displacement during binning and uses the full Euclidean distance for proximity.

Parameters:

  • state (State) – Simulation state (positions, radii, etc.).

  • system (System) – System configuration including domain.

  • lidar_range (float) – Maximum detection range and reference distance for proximity.

  • n_bins (int) – Number of angular bins (rays) spanning \([-\pi, \pi)\).

  • sense_edges (bool, optional) – If True, the function includes domain boundaries as proximity sources. Wall detections receive an ID of -1. Only meaningful for bounded domains. Default is False.

Returns:

(state, system, proximity, ids, overflow) where state and system are unchanged, proximity and ids have shape (N, n_bins), and overflow is always False. Bins with no detection have ids set to the particle’s own index.

Return type:

Tuple[State, System, jax.Array, jax.Array, jax.Array]

Notes

This function computes all-pairs displacements directly from state.pos and does not invoke the collider. The returned ids are indices into state.pos in its current order.

Examples

>>> state, system, prox, ids, overflow = lidar_2d(state, system,
...     lidar_range=5.0, n_bins=36)
jaxdem.utils.environment.lidar_3d(state: State, system: System, lidar_range: float, n_azimuth: int, n_elevation: int, sense_edges: bool = False) → tuple[State, System, jax.Array, jax.Array, jax.Array][source]#

3-D LIDAR proximity readings and neighbor IDs.

Similar to lidar_2d() but bins neighbors on a spherical grid defined by n_azimuth azimuthal sectors in \([-\pi, \pi)\) and n_elevation elevation bands in \([-\pi/2, \pi/2]\). The returned proximity and ID arrays have shape (N, n_azimuth * n_elevation) with flat indexing az * n_elevation + el.

Parameters:

  • state (State) – Simulation state.

  • system (System) – System configuration including domain.

  • lidar_range (float) – Maximum detection range and reference distance for proximity.

  • n_azimuth (int) – Number of azimuthal bins.

  • n_elevation (int) – Number of elevation bins.

  • sense_edges (bool, optional) – If True, the function includes domain boundaries as proximity sources. Wall detections receive an ID of -1. Default is False.

Returns:

(state, system, proximity, ids, overflow) where state and system are unchanged, proximity and ids have shape (N, n_azimuth * n_elevation), and overflow is always False.

Return type:

Tuple[State, System, jax.Array, jax.Array, jax.Array]

Notes

Uses an all-pairs approach and does not invoke the collider. Returned ids index into state.pos in its current order.

Examples

>>> state, system, prox, ids, overflow = lidar_3d(state, system,
...     lidar_range=5.0, n_azimuth=36, n_elevation=18)
jaxdem.utils.environment.cross_lidar_2d(pos_a: jax.Array, pos_b: jax.Array, system: System, lidar_range: float, n_bins: int) → tuple[jax.Array, jax.Array, jax.Array][source]#

2-D LIDAR proximity and IDs from pos_a sensing targets in pos_b.

Computes all-pairs displacements from pos_a to pos_b, bins by azimuthal angle, and returns per-bin proximity and closest target IDs.

Parameters:

  • pos_a (jax.Array) – Sensor positions, shape (N_A, dim).

  • pos_b (jax.Array) – Target positions, shape (N_B, dim).

  • system (System) – System configuration.

  • lidar_range (float) – Maximum detection range and reference distance for proximity.

  • n_bins (int) – Number of angular bins spanning \([-\pi, \pi)\).

Returns:

(proximity, ids, overflow) where proximity and ids have shape (N_A, n_bins) and overflow is always False. Empty bins get ids = -1.

Return type:

Tuple[jax.Array, jax.Array, jax.Array]

Notes

Uses an all-pairs approach and does not invoke the collider. Returned ids are indices into pos_b.

Examples

>>> prox, ids, overflow = cross_lidar_2d(agents, obstacles, system,
...                                      lidar_range=5.0, n_bins=36)
jaxdem.utils.environment.cross_lidar_3d(pos_a: jax.Array, pos_b: jax.Array, system: System, lidar_range: float, n_azimuth: int, n_elevation: int) → tuple[jax.Array, jax.Array, jax.Array][source]#

3-D LIDAR proximity and IDs from pos_a sensing targets in pos_b.

Computes all-pairs displacements from pos_a to pos_b, bins on a spherical grid, and returns per-bin proximity and closest target IDs.

Parameters:

  • pos_a (jax.Array) – Sensor positions, shape (N_A, 3).

  • pos_b (jax.Array) – Target positions, shape (N_B, 3).

  • system (System) – System configuration.

  • lidar_range (float) – Maximum detection range and reference distance for proximity.

  • n_azimuth (int) – Number of azimuthal bins.

  • n_elevation (int) – Number of elevation bins.

Returns:

(proximity, ids, overflow) where proximity and ids have shape (N_A, n_azimuth * n_elevation) and overflow is always False. Empty bins get ids = -1.

Return type:

Tuple[jax.Array, jax.Array, jax.Array]

Notes

Uses an all-pairs approach and does not invoke the collider. Returned ids are indices into pos_b.

Examples

>>> prox, ids, overflow = cross_lidar_3d(agents, obstacles, system,
...                                      lidar_range=5.0, n_azimuth=36,
...                                      n_elevation=18)