Introduce a Custom Environment

TIP

Install with pip install telekinesis-rlbotics — see Install Support for Reinforcement Learning for the simulator backends.

SUMMARY

Nothing in the library is tied to a simulator. Implement VecEnv — three methods and a handful of attributes — and the runner trains against it.

For a humanoid or quadruped already in mjlab or Isaac Lab there is no VecEnv to write: subclass the adapter instead, and override only what your task needs.

Start from this stub

Copy it, fill the blocks, delete the NotImplementedErrors as you go.

import torch
from tensordict import TensorDict

from telekinesis.rlbotics.envs.base import VecEnv


class MyEnv(VecEnv):
    """One environment stepped in parallel across num_envs copies."""

    def __init__(self, num_envs: int = 64, device: str = "cpu") -> None:
        # --- required by the runner ---
        self.num_envs = num_envs
        self.num_actions = ...              # fill: usually the joint count
        self.device = torch.device(device)
        self.max_episode_length = ...       # fill: steps before the time limit cuts an episode
        self.cfg = {}                       # fill: your own config, for logging

        # --- optional ---
        # Set these only if actions really are bounded (e.g. joint limits). export() then bakes
        # the clipping and scaling into the ONNX graph. Leave them off for raw joint targets.
        # self.action_low = torch.full((self.num_actions,), -1.0, device=self.device)
        # self.action_high = torch.full((self.num_actions,), 1.0, device=self.device)

        # Needed only for learn(init_at_random_ep_len=True)
        self.episode_length_buf = torch.zeros(num_envs, dtype=torch.long, device=self.device)

        self.reset()

    def _observations(self) -> TensorDict:
        """Build the observation groups. The names here are what obs_groups refers to."""
        raise NotImplementedError(
            # fill: one entry per group, each shaped (num_envs, ...)
            # return TensorDict(
            #     {
            #         "policy": ...,        # what the robot can measure
            #         "privileged": ...,    # optional: what only the simulator knows
            #     },
            #     batch_size=(self.num_envs,),
            #     device=self.device,
            # )
        )

    def reset(self) -> TensorDict:
        """Reset every environment and return the observation they start from."""
        raise NotImplementedError  # fill

    def get_observations(self) -> TensorDict:
        """Return the current observations without stepping."""
        return self._observations()

    def step(self, actions: torch.Tensor):
        """Apply one batch of actions and report the outcome."""
        # 1. Advance the simulation
        raise NotImplementedError  # fill

        # 2. Reward and episode end
        # rewards = ...                                  # (num_envs,)
        # terminated = ...                               # (num_envs,) bool: the task failed
        # time_outs = self.episode_length_buf >= self.max_episode_length
        # dones = terminated | time_outs

        # 3. Reset the finished environments now, so every step returned is a real transition
        # finished = dones.nonzero(as_tuple=False).squeeze(-1)
        # if finished.numel() > 0:
        #     self._respawn(finished)

        # 4. time_outs kept separate from dones, so a truncated episode is bootstrapped
        # return self._observations(), rewards, dones, {"time_outs": time_outs}

Train it:

from telekinesis.rlbotics.config import MLPConfig, OnPolicyRunnerConfig
from telekinesis.rlbotics.runner import create_runner

env = MyEnv(num_envs=64)
cfg = OnPolicyRunnerConfig(
    obs_groups={"actor": ["policy"], "critic": ["policy", "privileged"]},
    num_learning_iterations=300,
    seed=0,                 # reproducible while you are debugging
    check_for_nan=True,     # catches a bad reward or observation before it reaches the models
)
runner = create_runner(env=env, runner_cfg=cfg, device="cpu")
runner.learn()
runner.export()

What the runner reads

Attribute	Required	What it is
num_envs	yes	Environments stepped in parallel. 1 for a real robot
num_actions	yes	Size of the action vector one environment expects
device	yes	Device the observations, rewards and dones are placed on
max_episode_length	yes	Steps before the time limit cuts an episode off
cfg	yes	Your own configuration, for logging and introspection
episode_length_buf	optional	Steps each environment is into its episode
action_low / action_high	optional	Action bounds. Set them and export() folds the scaling into the graph
render() / render_fps	optional	Video recording. Without them log_video: true warns and writes nothing
close() / seed(seed)	optional	Cleanup, and seeding the simulation — runner.seed covers only the networks and sampling

VecEnv is an interface, not a base class that holds state: it has no __init__, and these are declarations of what gets read. That is deliberate — an environment usually does not own these numbers, it reads them off the simulator it wraps.

Warnings

Reset inside step

Every step the runner receives has to be a real transition, so finish an episode and reset it in the same call, returning the observation the new episode starts from.

The bootstrap value is unaffected: the algorithm uses the value it computed from the observation it acted on, before the step.

Keep time_outs out of the failure signal

dones is True for both reasons — the robot fell over, or the clock ran out — while extras["time_outs"] is True for the time limit alone.

Without the split, the policy is taught that surviving to the limit is as bad as falling over. The run still trains; it just trains worse, which is why this one is easy to miss.

mjlab Environment

mjlab combines Isaac Lab's manager-based API with MuJoCo Warp. If your robot is already a registered mjlab task, MjlabVecEnv reads everything off it — num_actions from action_manager.total_action_dim, the observation groups from observation_manager.compute(), the episode limit from the task config — and there is nothing to write.

Subclass it when you need to change how the task is presented:

import torch
from tensordict import TensorDict

from telekinesis.rlbotics.envs.mjlab_env import MjlabVecEnv


class MyMjlabEnv(MjlabVecEnv):
    """An mjlab task with something added on top."""

    def __init__(self, task: str, num_envs: int = 4096, device: str = "auto", **kwargs) -> None:
        super().__init__(task=task, num_envs=num_envs, device=device, **kwargs)
        # fill: anything your task needs on top, e.g. a command curriculum or buffers of your own
        raise NotImplementedError

    def _observations(self, obs: dict[str, torch.Tensor]) -> TensorDict:
        """Reshape or extend what the task publishes."""
        observations = super()._observations(obs)
        # fill: e.g. add a group of your own, or concatenate two the task publishes separately
        # observations["privileged"] = ...
        raise NotImplementedError

    def step(self, actions: torch.Tensor):
        """Wrap a step — extra logging, action shaping, a custom termination."""
        obs, rewards, dones, extras = super().step(actions)
        # fill: e.g. add a reward term, or extras["log"] entries
        raise NotImplementedError

Authoring a new task is mjlab's own job rather than RLBotics': you configure a scene, an ObservationManager with named groups, an ActionManager, a RewardManager, a TerminationManager and an EventManager for randomisation. See mujocolab.github.io/mjlab. Once registered, point a config at it:

env:
  framework: mjlab
  id: My-Humanoid-Task
  num_envs: 4096
  device: auto
  clip_actions: 1.0
  # nconmax: 50000        # raise on "nconmax overflow" in a contact-heavy scene

Truncations only reach the algorithm on an infinite-horizon task

The adapter forwards them as extras["time_outs"] when cfg.is_finite_horizon is false. On a finite-horizon task the limit is part of the task, so no bootstrap happens.

That is correct for a task with a real deadline, and wrong if you meant a locomotion task to run indefinitely — check which one your task config declares.

Actions are joint targets, so there are no action bounds

The adapter exposes none, which means clipping is env.clip_actions and an exported policy carries the observation normalization but no action scaling. Your deployment applies the same clip — see action scaling.

Isaac Lab Environment

Same managers, on Isaac Sim. Before writing anything, check the Available Environments — the locomotion set covers ANYmal B/C/D, Unitree A1/Go1/Go2, Spot, and the H1, G1 and Digit humanoids on flat and rough terrain, and RLBotics ships 31 configs across them.

import torch
from tensordict import TensorDict

from telekinesis.rlbotics.envs.isaaclab_env import IsaacLabVecEnv


class MyIsaacLabEnv(IsaacLabVecEnv):
    """An Isaac Lab task with something added on top."""

    def __init__(self, task: str, num_envs: int = 4096, device: str = "auto", **kwargs) -> None:
        # Isaac Sim is launched inside here, before the task is imported
        super().__init__(task=task, num_envs=num_envs, device=device, **kwargs)
        # fill
        raise NotImplementedError

    def _observations(self, obs: dict[str, torch.Tensor]) -> TensorDict:
        """Reshape or extend what the task publishes."""
        observations = super()._observations(obs)
        # fill
        raise NotImplementedError

    def step(self, actions: torch.Tensor):
        obs, rewards, dones, extras = super().step(actions)
        # fill
        raise NotImplementedError

For a new task, Isaac Lab's Creating a Manager-Based RL Environment is the reference: ManagerBasedRLEnvCfg with ObservationGroupCfg / ObservationTermCfg, RewardTermCfg, TerminationTermCfg and EventTermCfg.

env:
  framework: isaaclab
  id: My-Humanoid-Task-v0
  num_envs: 4096
  device: auto
  headless: true

Import order

Isaac Sim has to be running before any Isaac Lab task is imported, and the adapter launches it in its own constructor. Nothing from isaaclab_tasks may be imported above that — including your own task module, if it imports from there.

Naming the observation groups

Whichever path you took, the group names are the vocabulary obs_groups uses:

Backend	Groups published
Gymnasium	observation
mjlab	actor, plus a privileged critic, plus camera on the vision tasks
Isaac Lab	policy, plus critic when the task is set up asymmetrically
Yours	Whatever you put in the TensorDict

For a legged robot the split is what the robot can measure versus what only the simulator knows.

Actor group — everything a real robot could produce on its own: base angular velocity from the IMU, projected gravity (orientation without an absolute yaw, so the policy does not learn a compass), joint positions and velocities from the encoders, the previous action, and the commanded velocity. For the Unitree G1's 29 joints that is 99 dimensions: base_lin_vel[0:3], base_ang_vel[3:6], projected_gravity[6:9], joint_pos[9:38], joint_vel[38:67], actions[67:96], command[96:99].

Privileged critic group — true base velocity rather than an estimate, per-foot contact forces and air time, a terrain height scan around the robot, the disturbance forces you are applying as randomisation. Only the actor is exported, so none of this costs anything at deployment: why the critic gets its own set.

A set naming several groups concatenates them and requires each to be flat. A set naming one group keeps its shape, which is what a CNNConfig needs for camera observations.

A left-right symmetric robot has a free invariance

Telling PPO about it is worth real sample efficiency, and is what stops a policy settling into a limp. It needs a mirror function over the layout above, which only you can write — symmetry covers deriving one from the joint names and, more importantly, verifying it. Rough terrain is not mirror-symmetric in its dynamics even when the observation map is right.

Validate before a long run

from telekinesis.rlbotics.envs.base import observation_spec

env = MyEnv(num_envs=4)
observation_spec(env)                      # {"policy": (48,), "privileged": (235,)}

obs, rewards, dones, extras = env.step(torch.zeros(4, env.num_actions))
rewards.shape, dones.shape                 # both (4,)
extras["time_outs"].shape                  # (4,)

Turn check_for_nan off once the environment works

It checks every environment output for NaN, which is what you want on a first run — but each check forces a GPU-to-CPU sync, so leaving it on noticeably caps utilization on a GPU-resident environment.

Wrapping a real robot

The same interface describes hardware: num_envs=1, max_episode_length from your safety timeout, observations assembled from your sensors in a fixed order, and step writing joint targets to the controller. Two things differ — reset means whatever your homing routine is, and the loop is wall-clock bound rather than compute bound, so num_steps_per_env becomes a real-time budget.

Inference on hardware needs no VecEnv at all

VecEnv is for training against a robot. To only run a trained policy, the exported file takes numpy in and gives numpy out — see Deploy a Policy.

Need something other than PPO?

The protocol the runner drives, how to subclass PPO, and a worked SAC implementation.

Introduce a Custom Algorithm →