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)#

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.

The reward uses potential-based shaping with a proximity-gated kinetic-energy term:

\[\varphi_i(d, K) = \exp\!\left(-2 d^{\mathrm{eff}} - \frac{K}{\text{ke\_tau}}\,e^{-\text{ke\_gate} \cdot d^{\mathrm{eff}}}\right)\]

where \(d^{\mathrm{eff}} = \max(0, d - 0.5 r)\), \(d\) is the distance to the assigned objective, and \(K\) is the translational kinetic energy. ke_tau sets the overall strength of the KE penalty. ke_gate controls how sharply KE sensitivity falls off with distance. A larger ke_gate means KE only matters very close to the objective. The per-agent shaping credit is \(F_i = \varphi_i(d^{\mathrm{eff}}_t, K_t) - \varphi_i(d^{\mathrm{eff}}_{t-1}, K_{t-1})\).

Notes

The observation vector per agent is:

Feature	Size
Unit direction to objective	`dim`
Clamped displacement	`dim`
Velocity	`dim`
LiDAR proximity (normalized)	`n_lidar_rays`

For realistic training parameters, skip_frames = 50 gives a response rate of 200 Hz, so num_steps_epoch = 100 gives a horizon of 0.5 seconds.

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

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, ke_tau: float = 5.0, ke_gate: float = 4.0, near_goal_bonus: float = 0.1, lidar_range: float = 10.0, n_lidar_rays: int = 16) → MultiNavigator[source]#

Create a multi-agent navigator environment.

Parameters:

N (int) – Number of agents.
min_box_size (float) – Range for the random square domain side length sampled at each reset().
max_box_size (float) – Range for the random square domain 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.
ke_tau (float) – Overall strength of the KE term in the potential (larger = less important). See class docstring.
ke_gate (float) – Distance decay rate of KE sensitivity (larger = KE only matters very close to the goal). See class docstring.
near_goal_bonus (float) – Reward bonus applied when an agent is within one radius of its objective.
lidar_range (float) – Maximum detection range for the LiDAR sensor.
n_lidar_rays (int) – Number of angular LiDAR bins spanning \([-\pi, \pi)\).

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]#

Advance one step. Actions are forces. The step also applies drag -friction * vel.

Parameters:

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

Returns:

The updated environment state.

Return type:

Environment

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

Build per-agent observations.

Contents per agent#

Unit vector to objective (shape (dim,)) –> Direction
Clamped delta to objective (shape (dim,)) –> Local precision
Velocity (shape (dim,))
LiDAR proximity, normalized by lidar_range (shape (n_lidar_rays,))

returns:: Array of shape (N, 3 * dim + n_lidar_rays)
rtype:: jax.Array

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

Return the per-agent rewards.

Potential-based shaping with a proximity-gated KE term:

\[\varphi(d, K) = \exp\!\left(-2 d^{\mathrm{eff}} - \frac{K}{\text{ke\_tau}}\,e^{-\text{ke\_gate} \cdot d^{\mathrm{eff}}}\right)\]

The gate \(e^{-\text{ke\_gate} \cdot d^{\mathrm{eff}}}\) suppresses the KE term away from the objective, so fast motion is free until the agent is close. ke_tau sets the overall strength of the penalty.

Per-step reward:

\[\mathrm{rew}_t = \frac{F_t + w_{\text{near}} \cdot \mathbf{1}[d_t \le r]}{w_{\text{near}}}\]

where \(F_t = \varphi(d^{\mathrm{eff}}_t, K_t) - \varphi(d^{\mathrm{eff}}_{t-1}, K_{t-1})\), \(d^{\mathrm{eff}}_t = \max(0, d_t - 0.5 r)\), and \(w_{\text{near}}\) weights a near-goal bonus.

Parameters:: env (Environment) – Current environment.
Returns:: Shape (N,).
Return type:: jax.Array

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

Return whether the episode has ended.

The episode ends when step_count exceeds 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).