jaxdem.rl.environments.single_roller#

Environment where a single agent rolls toward a target on the floor.

Functions

frictional_wall_force(pos, state, system)

Normal, frictional, and restitution forces for a sphere on a \(z = 0\) plane.

Classes

SingleRoller(state, system, env_params)

Single-agent 3D navigation through torque-controlled rolling.
class jaxdem.rl.environments.single_roller.SingleRoller(state: State, system: System, env_params: dict[str, Any])#

Bases: Environment

Single-agent 3D navigation through torque-controlled rolling.

The agent is a sphere resting on a \(z = 0\) floor under gravity. Actions are 3-D torque vectors. Translational motion comes from frictional contact with the floor (see frictional_wall_force()). Each step applies a viscous drag -friction * vel and an angular damping -friction * ang_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 total (translational + rotational) 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

2

Clamped displacement (x, y)

2

Velocity (x, y)

2

Angular velocity

3

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(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 = 5.0, ke_gate: float = 4.0) → SingleRoller[source]#

Create a single-agent roller environment.

Parameters:

  • 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) – Damping coefficient applied as -friction * vel and -friction * ang_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:

SingleRoller

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

Place the agent and the objective at random positions on the floor.

Parameters:

  • env (Environment) – The current environment.

  • key (ArrayLike) – JAX random number generator key.

Returns:

The initialized environment.

Return type:

Environment

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

Apply a torque action and advance the physics by one step.

Parameters:

  • env (Environment) – Current environment.

  • action (jax.Array) – 3-D torque vector per agent.

Returns:

Updated environment after one physics step.

Return type:

Environment

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

Per-agent observation vector.

Contents per agent:

  • Unit displacement to objective projected to x-y (shape (2,)).

  • Clamped displacement to objective projected to x-y (shape (2,)).

  • Velocity projected to x-y (shape (2,)).

  • Angular velocity (shape (3,)).

Returns:

Shape (N, 9).

Return type:

jax.Array

static reward(env: SingleRoller) → 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.

Returns:

Shape (N,).

Return type:

jax.Array

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

True when step_count exceeds max_steps.

property action_space_size: int[source]#

Per-agent flattened action dimensionality (3-D torque).

property action_space_shape: tuple[int][source]#

Per-agent action tensor shape.

property observation_space_size: int[source]#

Per-agent flattened observation dimensionality (9).