jaxdem.rl.environments.burrowing_dp#

Classes

BurrowingDP(state, system, env_params)

2D environment where a deformable particle burrows into a granular bed.

class jaxdem.rl.environments.burrowing_dp.BurrowingDP(state: State, system: System, env_params: dict[str, Any])#

Bases: Environment

2D environment where a deformable particle burrows into a granular bed.

The agent actuates internal cross-springs to squirm through the spheres.

classmethod Create(box_size: float = 40.0, max_steps: int = 10000) → BurrowingDP[source]#

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

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

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

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

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

property max_num_agents: int[source]#: Maximum number of active agents in the environment.

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).