jaxdem.rl.environments.single_navigator#

Environment where a single agent navigates toward a target.

Classes

SingleNavigator(state, system, env_params)

Single-agent navigation environment toward a fixed target.
class jaxdem.rl.environments.single_navigator.SingleNavigator(state: State, system: System, env_params: dict[str, Any])#

Bases: Environment

Single-agent navigation environment toward a fixed target.

The agent controls a force vector that acts directly on a sphere inside a reflective box. Each step adds viscous drag -friction * vel. The reward uses potential-based shaping with a proximity-gated kinetic-energy term:

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

where \(d\) is the distance to the objective and \(K\) is the translational kinetic energy. ke_tau is the KE scale that sets the overall strength of the 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 shaping credit is \(F_t = \varphi(d_t, K_t) - \varphi(d_{t-1}, K_{t-1})\), so kinetic energy is penalized only near the objective. Far away the gate \(e^{-\text{ke\_gate} \cdot d} \to 0\) and fast motion is free.

Per-step reward:

\[\mathrm{rew}_t = \frac{F_t + b \cdot \mathbb{1}[d_t \le r]}{b}\]

where \(b\) is the near-goal bonus and \(r\) is the agent radius.

Notes

The observation vector per agent is:

Feature

Size

Unit direction to objective

dim

Clamped displacement

dim

Velocity

dim

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.

classmethod Create(dim: int = 2, min_box_size: float = 40.0, max_box_size: float = 40.0, max_steps: int = 20000, friction: float = 0.2, near_goal_bonus: float = 0.1, ke_tau: float = 2.0, ke_gate: float = 6.0) → SingleNavigator[source]#

Create a single-agent navigator environment.

Parameters:

  • dim (int) – Spatial dimensionality (2 or 3).

  • min_box_size (float) – Range for the random square domain side length.

  • max_box_size (float) – Range for the random square domain side length.

  • max_steps (int) – Episode length in physics steps.

  • friction (float) – Viscous drag coefficient applied as -friction * vel.

  • near_goal_bonus (float) – Reward bonus applied when the agent is within one radius of the objective.

  • 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.

Returns:

The constructed environment. Call reset() before use.

Return type:

SingleNavigator

static reset(env: SingleNavigator, key: Array | ndarray | bool | number | bool | int | float | complex) → Environment[source]#

Place the agent and the objective at random positions in the box.

Parameters:

  • env ('SingleNavigator') – The current environment.

  • key (jax.random.PRNGKey) – JAX random number generator key.

Returns:

The initialized environment.

Return type:

Environment

static step(env: SingleNavigator, 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: SingleNavigator) → 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,))

returns:

Array of shape (N, 3 * dim)

rtype:

jax.Array

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

Return the per-agent rewards.

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

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

The gate \(e^{-\text{ke\_gate} \cdot d}\) 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{\varphi(d_t, K_t) - \varphi(d_{t-1}, K_{t-1}) + b \cdot \mathbb{1}[d_t \le r]}{b}\]

where \(b\) is the near-goal bonus and \(r\) is the agent radius.

Parameters:

env (Environment) – Current environment.

Returns:

Shape (N,).

Return type:

jax.Array

static done(env: SingleNavigator) → 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).