torch_world.envs¶
PointMass(
batch_size: int = 1,
size: int = 32,
max_step: float = 0.1,
sigma: float = 0.06,
device=None,
)
Batched 2D point mass in the unit square, rendered as a Gaussian blob on an size x size RGB image.
Actions are velocities in [-1, 1]^2 scaled by max_step; positions are clipped to [low, high]^2.
sample_goals(
distance: tuple[float, float] = (0.2, 0.5),
generator: Generator | None = None,
) -> Tensor
Goal states at a uniform distance in distance from the current states, inside the bounds.
Images (B, 3, size, size) in [0, 1]: a dark blob on a light color gradient.
staticmethod
¶
The planar position to reach: the point itself.
Distance to goal (B, state_dim) of states (B, ..., state_dim).
Random-action episodes: pixels (B, T, 3, size, size), actions (B, T, 2), states (B, T, 2).
actions[:, t] is applied at states[:, t].
Macaw(
batch_size: int = 1,
size: int = 64,
max_step: float = 0.1,
scale: float = 1.0,
device=None,
)
Bases: PointMass
A Hyacinth macaw flying over the Pantanal, on size x size RGB images: the same dynamics as PointMass.
The state is the bird's position, kept in [low, high]^2. Its wings beat a quarter cycle per step, so a single
frame does not tell the phase but a short history does. render(state) without a phase draws the bird gliding,
as used for goal images. Rendering is resolution-free: a Macaw(size=192) draws the same states sharper.
Images (B, 3, size, size) in [0, 1]; phase (B,) sets the wing beat.
RobotCube(
robot: str = "ur10e",
batch_size: int = 1,
size: int = 64,
camera: str = "front",
control_dt: float = 0.1,
max_step: float | None = None,
task: str = "pick",
root: str | Path | None = None,
device=None,
)
Bases: MujocoEnv
Single-cube pick-and-place with a real arm in MuJoCo, after OGBench cube-single, in a light studio scene.
robot is "ur10e" (Universal Robots UR10e with a Robotiq 2F-85 gripper) or "so101" (the open-source
SO-101 arm), both from MuJoCo Menagerie and downloaded once into root. Actions (dx, dy, dz, grip) in
[-1, 1]^4 move the gripper target by up to max_step meters, solved by damped least-squares inverse
kinematics with the gripper pointing down, and open (grip > 0) or close the gripper. Each step lasts
control_dt seconds of simulation. The state is (gripper x, y, z, opening, cube x, y, z).
task sets what tasks, distance, success and position measure: "pick" carries the cube to a
goal spot, "reach" moves the open gripper to a goal pose with the cube left in place.
The batch is simulated sequentially on the CPU; observations are (B, 3, size, size) renders of the
"front" camera ("top" is an orthographic view of the workspace, "wrist" rides on the gripper).
Park et al., "OGBench: Benchmarking Offline Goal-Conditioned RL", ICLR 2025. Paper: https://arxiv.org/abs/2410.20092 Code: https://github.com/seohongpark/ogbench Models: https://github.com/google-deepmind/mujoco_menagerie
A random cube on the table and the open gripper hovering above the workspace, or the given state.
Random starts, and goals the scripted expert reaches from them: the cube carried to a random spot with the
gripper backed off ("pick", run until every cube is placed, up to 60 steps), or the open gripper moved to a
random pose above the workspace ("reach").
Clean expert actions (B, 4) toward goal states (B, 7) of the current task.
The cube's ("pick") or the gripper's ("reach") position on the table in the unit square of the
"top" camera (x right, y down), which spans top_size meters.
Position in meters of what the task moves to its goal, for path-based metrics such as SPL: the cube on the
table plane ("pick"; lifting it is not a detour) or the gripper in 3D ("reach").
Distance in meters between the cube and its goal spot on the table ("pick"), or between the gripper and
its goal position ("reach"), for states (B, ..., 7) and goals (B, 7).
The cube rests on the table within radius of its goal ("pick"), or the gripper is within radius
of its goal ("reach").
Goal cube positions (B, 2) inside the workspace, at least distance[0] and at most distance[1] away.
Scripted pick-and-place actions (B, 4) that carry the cube to target (B, 2) and back off.
Motion is planned from the commanded gripper position of the running episode, or from state.
Run the clean expert for steps; returns the final states, used as goals.
collect(
num_steps: int,
noise: float = 0.04,
cameras: list[str] | None = None,
generator: Generator | None = None,
) -> dict
Play data from the noisy expert, retargeted after every placement, as in OGBench's play datasets.
Returns pixels (B, T, 3, size, size) (or (B, T, V, 3, size, size) for a list of cameras),
actions (B, T, 4) and states (B, T, 7); actions[:, t] is applied at states[:, t].
noise is the standard deviation added to the expert's actions.
Run policy(observation) -> action in closed loop from the current state; returns states (B, steps + 1, ...).