jaxdem.rl.environments.multi_navigator#

Environment where multiple agents navigate toward assigned targets.

Classes

MultiNavigator(state, system, env_params, ...)

Multi-agent navigation environment toward assigned targets.

class jaxdem.rl.environments.multi_navigator.MultiNavigator(state: State, system: System, env_params: dict[str, Any], n_lidar_rays: int, n_lidar_elevation: int)#

Bases: Environment

Multi-agent navigation environment toward assigned targets.

Each agent controls a force vector that acts directly on a sphere inside a reflective box. Each step adds viscous drag -friction * vel. The environment samples objectives and assigns them one-to-one with a random permutation.

Each physics step smooths each requested force:

\[\mathbf{u}_t = \alpha\,\mathbf{a}_t + (1 - \alpha)\,\mathbf{u}_{t-1}\]

Here \(\mathbf{a}_t\) is the requested force at physics step \(t\), \(\mathbf{u}_t\) is the applied force before drag, and \(\mathbf{u}_{t-1}\) is the immediately preceding applied force. The parameter \(\alpha\) is action_alpha. Reset initializes \(\mathbf{u}_0 = \mathbf{0}\); checkpoints preserve this actuator state. With \(\alpha = 0.22\), a constant request completes 99% of its transition in 19 physics steps.

The reward uses a quartic exponential potential measured from the action-start checkpoint to the live endpoint:

\[R_i = 2\,\mathrm{rad}_i, \qquad K_i = \tfrac12 m_i \|\mathbf{v}_i\|^2, \qquad \Phi_i(d,K) = \exp\!\left[-\left(\frac{d}{R_i}\right)^4 - \beta\frac{K}{K_{\mathrm{ref}}}\right].\]

Here \(d\) is the distance from agent \(i\)’s center to its assigned objective, and \(\mathrm{rad}_i\) is its particle radius. The potential is flat near the objective, making small departures inexpensive, and decays rapidly far away. Its scale follows each particle’s own radius. kinetic_energy_coeff sets \(\beta\) (default zero), and kinetic_energy_scale sets \(K_{\mathrm{ref}}\) (default one). Rotation is disabled, so only translational energy enters. At fixed distance, slowing down earns credit near the objective and has little effect far away. This is a potential difference, not a per-step cost for maintaining constant speed.

Per-action reward:

\[r_{i,t} = \Phi_i(d_{i,t},K_{i,t}) - \Phi_i(d_{i,\mathrm{checkpoint}},K_{i,\mathrm{checkpoint}}).\]

At fixed energy, moving closer gives positive credit and moving farther gives negative credit. Unchanged distance and energy give zero. Reset initializes the historical potential from the initial state, giving zero reward. There are no additional neighbor, occupancy, or collision reward terms.

Notes

The observation vector per agent is:

Feature	Size
Unit direction to objective	`dim`
Clamped displacement	`dim`
Velocity	`dim`
Normalized agent/wall LiDAR	`n_lidar_rays` in 2-D; `n_lidar_rays * n_lidar_elevation` in 3-D

With dt = 0.002 and skip_frames = 49, each action spans 50 physics steps (0.1 seconds). One checkpoint captures each agent’s prev_potential before the action; reward uses the live endpoint after all accepted physics steps. prev_dist retains the starting distance for inspection.

n_lidar_rays: int#: Number of angular bins for each LiDAR sensor.

n_lidar_elevation: int#: Number of elevation bins in 3-D; unused in 2-D.

classmethod Create(N: int = 64, min_box_size: float = 20.0, max_box_size: float = 20.0, box_padding: float = 5.0, max_steps: int = 100000, friction: float = 0.2, action_alpha: float = 0.22, lidar_range: float = 10.0, n_lidar_rays: int = 16, kinetic_energy_coeff: float = 0.0, kinetic_energy_scale: float = 1.0, dim: int = 2, n_lidar_elevation: int = 4) → MultiNavigator[source]#

Create a multi-agent navigator environment.

Parameters:

N (int) – Number of agents.
min_box_size (float) – Range for the random square or cube side length sampled at each reset().
max_box_size (float) – Range for the random square or cube side length sampled at each reset().
box_padding (float) – Extra padding around the domain in multiples of the particle radius.
max_steps (int) – Episode length in physics steps.
friction (float) – Viscous drag coefficient applied as -friction * vel.
action_alpha (float) – Fraction of the requested force applied by the smoothing update each physics step, in [0, 1]. One disables smoothing.
lidar_range (float) – Maximum detection range for the LiDAR sensor.
n_lidar_rays (int) – Number of angular LiDAR bins spanning \([-\pi, \pi)\).
kinetic_energy_coeff (float) – Nonnegative dimensionless strength of the kinetic-energy factor. Zero preserves the distance-only potential.
kinetic_energy_scale (float) – Positive reference kinetic energy in simulation units. The energy factor is exp(-kinetic_energy_coeff * K / kinetic_energy_scale).
dim (int) – Spatial dimension, either 2 (default) or 3.
n_lidar_elevation (int) – Number of elevation bins for 3-D LiDAR; unused in 2-D.

Returns:

The constructed environment. Call reset() before use.

Return type:

MultiNavigator

Initialize the environment with random positions and objectives.

Parameters:

env (Environment) – The current environment.
key (ArrayLike) – JAX random number generator key.

Returns:

The initialized environment.

Return type:

Environment

static step(env: MultiNavigator, action: Array) → Environment[source]#

Smooth the requested force, then advance physics with viscous drag.

Smoothing uses the applied action from the immediately preceding physics step, including across policy actions and checkpoints.

The historical potential remains fixed throughout the action. Use utils.advance_action to checkpoint once before repeating physics steps. Observations and rewards always read the live state.

Parameters:

env (Environment) – The current environment.
action (jax.Array) – The per-agent action vectors.

Returns:

The updated environment state.

Return type:

Environment

static checkpoint(env: MultiNavigator, action: Array) → Environment[source]#

Save the starting potential before the next action interval.

Store the live distance in prev_dist and the full distance/energy potential in prev_potential. Physics steps preserve both; reward subtracts the saved potential from the live endpoint’s potential.

Parameters:

env (MultiNavigator) – The environment immediately before the next action.
action (jax.Array) – The per-agent action about to be applied. This potential baseline does not depend on its value.

Returns:

The environment with the action-start potential saved. Read reward before the next checkpoint or reset replaces this baseline.

Return type:

Environment

static observation(env: MultiNavigator) → Array[source]#

Build per-agent observations and LiDAR from the live physical state.

Contents per agent#

Unit vector from objective to agent (shape (dim,)) –> Direction
Clamped agent-minus-objective delta (shape (dim,)) –> Local precision
Velocity (shape (dim,))
Agent/wall LiDAR proximity (shape (n_lidar_rays,) in 2-D, (n_lidar_rays * n_lidar_elevation,) in 3-D)

LiDAR reports normalized proximity \(p_k = \max(0, 1 - d_{\min,k} / r_{\max})\), with zero for empty bins. Readings use live positions rather than checkpointed values.

returns:: Array of shape (N, observation_space_size)
rtype:: jax.Array

static reward(env: MultiNavigator) → Array[source]#

Return the potential change over the complete action interval.

\[\Phi_i(d,K) = \exp\!\left[-\left( \frac{d}{2\,\mathrm{rad}_i}\right)^4 - \beta\frac{K}{K_{\mathrm{ref}}}\right], \qquad r_{i,t} = \Phi_i(d_{i,t},K_{i,t}) - \Phi_i(d_{i,\mathrm{checkpoint}},K_{i,\mathrm{checkpoint}}).\]

kinetic_energy_coeff is \(\beta\) and kinetic_energy_scale is \(K_{\mathrm{ref}}\). Energy is translational because navigator rotation is disabled. With the default zero coefficient, this is the distance-only quartic potential. The complete starting potential is saved by reset/checkpoint; physics steps preserve that baseline. Returns per-agent rewards of shape (N,).

static truncated(env: MultiNavigator) → Array[source]#

Return whether the episode has ended.

The episode ends when step_count reaches max_steps.

Parameters:: env (Environment) – The current environment.
Returns:: A bool that is True when the episode has ended.
Return type:: jax.Array

property action_space_size: int[source]#: Flattened action size per agent. Actions passed to step() have shape (A, action_space_size).

property action_space_shape: tuple[int][source]#: Original per-agent action shape (useful for reshaping inside the environment).

property observation_space_size: int[source]#: Flattened observation size per agent. observation() returns shape (A, observation_space_size).