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. Each physics step smooths the 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 an exponential distance potential measured from the action-start checkpoint to the live endpoint:

\[\varphi(d) = \exp\!\left(-2 d \right)\]

where \(d\) is the distance to the objective.

The shaping credit is \(F_t = \varphi(d_t) - \varphi(d_{t-1})\), which is positive when moving closer, negative when moving farther away, and zero when the distance is unchanged. Reset initializes the historical distance from the initial physical state, giving zero reward.

Per-action reward:

\[\mathrm{rew}_t = F_t\]

Notes

The observation vector per agent is:

Feature	Size
Unit direction to objective	`dim`
Clamped displacement	`dim`
Velocity	`dim`

With dt = 0.002 and skip_frames = 49, each action spans 50 physics steps (0.1 seconds). One checkpoint captures the distance before the action; the reward uses the live distance after all accepted physics steps. Only prev_dist is stored as reward history.

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, action_alpha: float = 0.22) → 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.
action_alpha (float) – Fraction of the requested force applied by the smoothing update each physics step, in [0, 1]. One disables smoothing.

Returns:

The constructed environment. Call reset() before use.

Return type:

SingleNavigator

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

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 distance 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: SingleNavigator, action: Array) → Environment[source]#

Save the starting distance before the next action interval.

Store the live distance in prev_dist. Physics steps preserve it, and reward compares it with the live endpoint after all repeated physics steps. Current quantities are computed directly from state.

Parameters:

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

Returns:

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

Return type:

Environment

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

Build per-agent observations directly 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,))

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

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

Return the distance-potential change since the action-start checkpoint.

\[\mathrm{rew}_t = e^{-2d_t} - e^{-2d_{t-1}}\]

Here \(d_t\) is the live distance from the agent’s center to its objective, and \(d_{t-1}\) is the distance saved before the action. The reward is positive for progress toward the objective and zero after reset or when the distance is unchanged. This accessor reads the live state and saved baseline without advancing the checkpoint.

Parameters:: env (Environment) – The live environment with its action-start distance baseline.
Returns:: Per-agent rewards of shape (N,), where N = 1.
Return type:: jax.Array

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