jaxdem.rl.environments.single_navigator#

Environment where a single agent navigates towards 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])[source]#

Bases: Environment

Single-agent navigation environment toward a fixed target.

classmethod Create(dim: int = 2, min_box_size: float = 1.0, max_box_size: float = 1.0, max_steps: int = 2000, final_reward: float = 2.0, shaping_factor: float = 0.0, prev_shaping_factor: float = 0.0, goal_threshold: float = 0.6666666666666666) → SingleNavigator[source][source]#

Custom factory method for this environment.

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

Initialize the environment with a randomly placed particle and velocity.

Parameters:

  • env (Environment) – Current environment instance.

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

Returns:

Freshly initialized environment.

Return type:

Environment

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

Advance one step. Actions are forces; simple drag is applied.

Parameters:

  • env (Environment) – The current environment.

  • action (jax.Array) – The vector of actions each agent in the environment should take.

Returns:

The updated environment state.

Return type:

Environment

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

Build per-agent observations.

Contents per agent#

  • Wrapped displacement to objective Δx (shape (2,)).

  • Velocity v (shape (2,)).

returns:

Array of shape (N, 2 * dim) scaled by the maximum box size for normalization.

rtype:

jax.Array

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

Returns a vector of per-agent rewards.

Equation

Let \(\delta_i=\operatorname{displacement}(\mathbf{x}_i,\mathbf{objective})\), \(d_i=\lVert\delta_i\rVert_2\), and \(\mathbf{1}[\cdot]\) the indicator. With shaping factors \(\alpha_{\text{prev}},\alpha\), final reward \(R_f\), and radius \(r_i\):

\[\mathrm{rew}^{\text{shape}}_i = \alpha_{\text{prev}}\,d^{\text{prev}}_i - \alpha\, d_i\]

Final reward:

\[\mathrm{rew}_i = \mathrm{rew}^{\text{shape}}_i + R_f\,\mathbf{1}[\,d_i < \text{goal_threshold}\times r_i\,]\]
Parameters:

env (Environment) – Current environment.

Returns:

Shape (N,). The normalized per-agent reward vector.

Return type:

jax.Array

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

Returns a boolean indicating whether the environment has ended. The episode terminates when the maximum number of steps is reached.

Parameters:

env (Environment) – The current environment.

Returns:

Boolean array indicating whether 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).