Skip to main content

Create an Environment

Create environments in your provider package when an installed environment no longer fits your scenario. The first task can reuse an existing environment; environment authoring is a separate step.

A minimal scene​

Create src/my_magma_scenarios/envs/__init__.py and this module:

src/my_magma_scenarios/envs/cube.py
import torch
from mani_skill.utils.building import actors
from mani_skill.utils.registration import register_env
from mani_skill.utils.scene_builder.table import TableSceneBuilder
from mani_skill.utils.structs import Pose
from magma_core.simulation.envs import DefaultEnv


@register_env("MyCube-v1", max_episode_steps=100)
class CubeEnv(DefaultEnv):
def __init__(self, *args, robot_uids="panda", **kwargs):
super().__init__(
*args, robot_uids=robot_uids, robot_init_qpos_noise=0.0, **kwargs
)

def _load_scene(self, options: dict) -> None:
self.table = TableSceneBuilder(self, robot_init_qpos_noise=0.0)
self.table.build()
self.cube = actors.build_cube(
self.scene, half_size=0.02, color=[0.2, 0.4, 0.8, 1], name="cube"
)
self.target = self.build_a_zone(width=0.1, length=0.1, name="target")

def _initialize_episode(self, env_idx: torch.Tensor, options: dict) -> None:
with torch.device(self.device):
self.table.initialize(env_idx)
xyz = torch.zeros((len(env_idx), 3))
xyz[:, 0] = -0.1
xyz[:, 2] = 0.02
self.cube.set_pose(Pose.create_from_pq(p=xyz))
target_xyz = torch.zeros((len(env_idx), 3))
target_xyz[:, 1] = 0.2
self.target.set_pose(Pose.create_from_pq(p=target_xyz))

def _get_obs_extra(self, info: dict) -> dict:
return {
"agent_tcp": self.agent.tcp.pose.raw_pose,
"cube": self.cube.pose.raw_pose,
"target": self.target.pose.raw_pose,
}

This module declares a scene, not a complete task. Its observation contract contains batched pose tensors under the three named keys. A task for this scene would use At("cube", "target") and tools able to manipulate the cube; button tools are not compatible.

Declare discovery​

Add to your manifest's environments mapping:

environments = {"MyCube-v1": "my_magma_scenarios.envs.cube"}

Set your cube task's maniskill_env_id = "MyCube-v1". The loader imports the environment module when loading a component from the scenario. list and show inspect declarations without constructing the scene.

You do not need to modify magma_scenarios/envs/__init__.py. For direct Gym use outside the MAGMA loader, import your registration module or call register_environment("MyCube-v1") before gym.make.

Preserve more than poses​

If you add machine states or counters, implement the extra-state hooks. Pose-only state is already handled by the simulator. Keep tensor batch dimensions compatible with restoring individual environment slots.

Assets and configuration​

Store assets and YAML in your own package and include them as package data:

Additional pyproject.toml section
[tool.setuptools.package-data]
my_magma_scenarios = ["assets/**/*", "scenarios/**/*.yaml"]

Use module-relative paths or package resources, not paths into a cloned MAGMA repository.

Next: Environment transitions. Environments explains ownership and resets.