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.

jaxdem.rl.environments.single_roller.frictional_wall_force(pos: Array, state: State, system: System) → tuple[Array, Array][source]#

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

Combines a linear spring in the normal direction with Coulomb tangential friction and a velocity-proportional dashpot for restitution damping.

Parameters:

pos (jax.Array) – Particle positions, shape (N, 3).
state (State) – Full simulation state (provides vel, ang_vel, rad, mass).
system (System) – System configuration (provides dt).

Returns:

total_force (jax.Array) – Per-particle force, shape (N, 3).
total_torque (jax.Array) – Per-particle torque, shape (N, 3).

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.

Each physics step smooths the requested torque:

\[\mathbf{u}_t = \alpha\,\mathbf{a}_t + (1 - \alpha)\,\mathbf{u}_{t-1}\]

Here \(\mathbf{a}_t\) is the requested torque at physics step \(t\), \(\mathbf{u}_t\) is the applied torque before damping, and \(\mathbf{u}_{t-1}\) is the immediately preceding applied torque. 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	2
Clamped displacement (x, y)	2
Velocity (x, y)	2
Angular velocity	3

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

Returns:

The constructed environment. Call reset() before use.

Return type:

SingleRoller

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

Smooth the requested torque and advance the physics by one step.

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) – Current environment.
action (jax.Array) – 3-D torque vector per agent.

Returns:

Updated environment after one physics step.

Return type:

Environment

static checkpoint(env: SingleRoller, 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 (SingleRoller) – 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: SingleRoller) → Array[source]#

Build per-agent observations directly from the live physical state.

Contents per agent:

Unit displacement from objective to agent projected to x-y (shape (2,)).
Clamped agent-minus-objective displacement 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 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:: Shape (N,).
Return type:: jax.Array

static truncated(env: SingleRoller) → Array[source]#: True when step_count reaches 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).