jaxdem.rl.environments.single_roller

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

Functions

frictional_wall_force(pos, state, system)

Calculates normal and frictional forces for a sphere on a y=0 plane.

Classes

SingleRoller3D(state, system, env_params)

Single-agent navigation where the agent rolls on a plane using torque control.

class jaxdem.rl.environments.single_roller.SingleRoller3D(state: State, system: System, env_params: Dict[str, Any])

Bases: Environment

Single-agent navigation where the agent rolls on a plane using torque control.

classmethod Create(min_box_size: float = 1.0, max_box_size: float = 1.0, max_steps: int = 6000, final_reward: float = 2.0, shaping_factor: float = 1.0, prev_shaping_factor: float = 1.0, goal_threshold: float = 0.6666666666666666) → SingleRoller3D

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

static step(env: Environment, action: Array) → Environment

static observation(env: Environment) → Array

static reward(env: Environment) → Array

static done(env: Environment) → Array

property action_space_size: int

Flattened action size per agent. Actions passed to step() have shape (A, action_space_size).

property action_space_shape: Tuple[int]

Original per-agent action shape (useful for reshaping inside the environment).

property observation_space_size: int

Flattened observation size per agent. observation() returns shape (A, observation_space_size).