Registering an Environment

Registering an Environment#

In the previous tutorial, we learned how to create a custom cartpole environment. We manually created an instance of the environment from the class named by its configuration’s class_type.

While straightforward, this approach is not scalable as we have a large suite of environments. In this tutorial, we will show how to use the gymnasium.register() method to register environments with the gymnasium registry. This allows us to create the environment through the gymnasium.make() function.

The Code#

The tutorial corresponds to the random_agent.py script in the scripts/environments directory. The script is a thin wrapper that calls into the isaaclab_rl.entrypoints module, where the actual implementation lives.

Code for simple_agents.py

# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Checkpoint-free playback workflows for Isaac Lab environments.

The zero and random agents are variations of playback that need no trained checkpoint:
the policy either infers finite zero or hold actions or samples uniform random actions.
"""

from __future__ import annotations

import argparse
import contextlib
import logging
import os
import sys
from collections.abc import Callable
from typing import Any, Literal

import gymnasium as gym
import torch

from isaaclab.app import add_launcher_args, launch_simulation
from isaaclab.envs.utils.spaces import sample_space
from isaaclab.utils import math as math_utils
from isaaclab.utils import validate

import isaaclab_tasks  # noqa: F401
from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli

from .common import (
    add_video_args,
    apply_env_overrides,
    apply_video_recording,
    close_env,
    enable_cameras_for_video,
    normalize_task_name,
    pre_launch_video_config,
    video_playback_steps,
)

logger = logging.getLogger(__name__)

# PLACEHOLDER: Extension template (do not remove this comment)
with contextlib.suppress(ImportError):
    import isaaclab_tasks_experimental  # noqa: F401

PolicyName = Literal["zero", "random"]
"""Action policies supported by the checkpoint-free agents."""

_DESCRIPTIONS: dict[str, str] = {
    "zero": "Zero agent for Isaac Lab environments.",
    "random": "Random agent for Isaac Lab environments.",
}

_SEED = 42


def run(argv: list[str] | None = None, *, policy: PolicyName) -> None:
    """Run an Isaac Lab environment with a checkpoint-free policy.

    Args:
        argv: Command-line arguments excluding the executable name. Reads ``sys.argv`` when omitted.
        policy: Action policy to apply, either inferred zero actions or uniform random actions.

    Raises:
        ValueError: If the requested policy is not supported.
    """
    if policy not in _DESCRIPTIONS:
        raise ValueError(f"Unsupported policy {policy!r}. Expected one of: {sorted(_DESCRIPTIONS)}.")

    args_cli = _parse_args(argv, policy)
    torch.manual_seed(_SEED)

    env_cfg, _ = resolve_task_config(args_cli.task, "")
    apply_env_overrides(args_cli, env_cfg)
    pre_launch_video_config(env_cfg, args_cli)
    # reject unsupported configurations before launching Kit or initializing a native physics backend
    try:
        validate(env_cfg)
    except (TypeError, ValueError) as exc:
        raise SystemExit(f"Invalid environment configuration: {exc}") from None

    with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup:
        log_dir = os.path.abspath(os.path.join("logs", f"{policy}_agent", normalize_task_name(args_cli.task)))
        apply_video_recording(env_cfg, log_dir, args_cli, subdir="play")
        env = gym.make(args_cli.task, cfg=env_cfg)
        cleanup.callback(lambda: close_env(env))
        logger.info(f"Gym observation space: {env.observation_space}")
        logger.info(f"Gym action space: {env.action_space}")
        env.reset()
        if policy == "zero":
            action_policy = create_zero_action_policy(env)
        else:
            action_policy = create_random_action_policy(env)
        logger.info(f"{policy.capitalize()} agent is running, press Ctrl+C to exit...")

        budgets = [n for n in (args_cli.max_steps, video_playback_steps(args_cli, env_cfg)) if n is not None]
        max_steps = min(budgets, default=None)

        # keep running while any visualizer is open and the step budget is not exhausted
        sim = env.unwrapped.sim
        step = 0
        while sim.is_running():
            if max_steps is not None and step >= max_steps:
                break
            step += 1
            with torch.inference_mode():
                env.step(action_policy())


def create_zero_action_policy(env: gym.Env) -> Callable[[], Any]:
    """Create a policy that emits finite actions for passive environment playback.

    Manager-based environments infer hold commands for absolute task-space action terms and use literal zeros for all
    other terms. Direct-workflow environments use zero-filled samples of their declared Gymnasium spaces, including
    composite and multi-agent spaces.
    """
    unwrapped = env.unwrapped
    action_manager = getattr(unwrapped, "action_manager", None)
    if action_manager is not None:
        return _create_manager_zero_action_policy(action_manager, unwrapped)

    if hasattr(unwrapped, "action_spaces"):
        actions = {
            agent: sample_space(space, unwrapped.device, batch_size=unwrapped.num_envs, fill_value=0)
            for agent, space in unwrapped.action_spaces.items()
        }
        return lambda: actions

    actions = sample_space(unwrapped.single_action_space, unwrapped.device, batch_size=unwrapped.num_envs, fill_value=0)
    return lambda: actions


def create_random_action_policy(env: gym.Env) -> Callable[[], torch.Tensor]:
    """Create a policy that samples uniform random actions in ``[-1, 1]``."""
    device = env.unwrapped.device
    return lambda: 2 * torch.rand(env.action_space.shape, device=device) - 1


def _create_manager_zero_action_policy(action_manager: Any, env: Any) -> Callable[[], torch.Tensor]:
    """Create a zero-action policy from the active action terms."""
    actions = torch.zeros_like(action_manager.action)
    term_policies = []
    index = 0
    for term_name in action_manager.active_terms:
        term = action_manager.get_term(term_name)
        term_policy = _create_action_term_zero_policy(term, env)
        if term_policy is not None:
            term_policies.append((slice(index, index + term.action_dim), term_policy))
        index += term.action_dim

    def policy() -> torch.Tensor:
        actions.zero_()
        for action_slice, term_policy in term_policies:
            actions[:, action_slice] = term_policy()
        if not torch.isfinite(actions).all():
            raise RuntimeError("Zero agent inferred non-finite actions from the current environment state.")
        return actions

    return policy


def _create_action_term_zero_policy(term: Any, env: Any) -> Callable[[], torch.Tensor] | None:
    """Create the specialized zero-action policy required by an action term."""
    term_types = {cls.__name__ for cls in type(term).__mro__}

    if "PinkInverseKinematicsAction" in term_types:
        controlled_frame_ids, controlled_frame_names = term._asset.find_bodies(
            list(term.cfg.target_eef_link_names.values()), preserve_order=True
        )
        if len(controlled_frame_ids) != len(term.cfg.target_eef_link_names):
            raise ValueError(
                "Expected one controlled body for every Pink IK target. Resolved "
                f"{controlled_frame_names} from {list(term.cfg.target_eef_link_names.values())}."
            )
        if len(controlled_frame_ids) != term._num_frame_tasks:
            raise ValueError(
                f"Pink IK has {term._num_frame_tasks} variable frame tasks but "
                f"{len(controlled_frame_ids)} controlled bodies were configured."
            )

        def pink_policy() -> torch.Tensor:
            frame_poses = term._asset.data.body_link_pose_w.torch[:, controlled_frame_ids].clone()
            frame_poses[..., :3] -= env.scene.env_origins.unsqueeze(1)
            hand_joint_positions = term._asset.data.joint_pos.torch[:, term._hand_joint_ids]
            return torch.cat((frame_poses.flatten(start_dim=1), hand_joint_positions), dim=-1)

        return pink_policy

    if "DifferentialInverseKinematicsAction" in term_types and not term.cfg.controller.use_relative_mode:

        def differential_ik_policy() -> torch.Tensor:
            ee_pos, ee_quat = term._compute_frame_pose()
            command = ee_pos if term.cfg.controller.command_type == "position" else torch.cat((ee_pos, ee_quat), dim=-1)
            return _unscale_action(command - term._offset, term._scale)

        return differential_ik_policy

    if "RMPFlowAction" in term_types and not term.cfg.use_relative_mode:

        def rmpflow_policy() -> torch.Tensor:
            ee_pos, ee_quat = term._compute_frame_pose()
            return _unscale_action(torch.cat((ee_pos, ee_quat), dim=-1), term._scale)

        return rmpflow_policy

    if "OperationalSpaceControllerAction" in term_types and term._pose_abs_idx is not None:
        term_actions = torch.zeros_like(term.raw_actions)

        def operational_space_policy() -> torch.Tensor:
            term_actions.zero_()
            term._compute_ee_pose()
            term._compute_task_frame_pose()
            if term._task_frame_pose_b is None:
                ee_pos_task = term._ee_pose_b[:, :3]
                ee_quat_task = term._ee_pose_b[:, 3:7]
            else:
                ee_pos_task, ee_quat_task = math_utils.subtract_frame_transforms(
                    term._task_frame_pose_b[:, :3],
                    term._task_frame_pose_b[:, 3:7],
                    term._ee_pose_b[:, :3],
                    term._ee_pose_b[:, 3:7],
                )
            term_actions[:, term._pose_abs_idx : term._pose_abs_idx + 3] = _unscale_action(
                ee_pos_task, term._position_scale
            )
            term_actions[:, term._pose_abs_idx + 3 : term._pose_abs_idx + 7] = _unscale_action(
                ee_quat_task, term._orientation_scale
            )
            return term_actions

        return operational_space_policy

    return None


def _unscale_action(command: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    """Map a processed command back to policy-action coordinates without division by zero."""
    return torch.where(scale != 0.0, command / scale, torch.zeros_like(command))


def _parse_args(argv: list[str] | None, policy: PolicyName) -> argparse.Namespace:
    """Parse the command line of a checkpoint-free agent and hand the remainder to Hydra."""
    parser = argparse.ArgumentParser(description=_DESCRIPTIONS[policy])
    parser.add_argument(
        "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
    )
    parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
    parser.add_argument("--task", type=str, default=None, help="Name of the task.")
    parser.add_argument(
        "--max_steps", type=int, default=None, help="Number of environment steps to run. Runs unbounded when omitted."
    )
    add_video_args(parser, action=f"the {policy} agent run")
    add_launcher_args(parser)
    # let task configs select the simulation device
    parser.set_defaults(device=None)
    args_cli, hydra_args = setup_preset_cli(parser, argv)
    enable_cameras_for_video(args_cli)
    sys.argv = [sys.argv[0]] + hydra_args
    return args_cli

The Code Explained#

The envs.ManagerBasedRLEnv class inherits from the gymnasium.Env class to follow a standard interface. However, unlike the traditional Gym environments, the envs.ManagerBasedRLEnv implements a vectorized environment. This means that multiple environment instances are running simultaneously in the same process, and all the data is returned in a batched fashion.

Similarly, the envs.DirectRLEnv class also inherits from the gymnasium.Env class for the direct workflow. For envs.DirectMARLEnv, although it does not inherit from Gymnasium, it can be registered and created in the same way.

Using the gym registry#

To register an environment, we use the gymnasium.register() method. This method takes in the environment name, the entry point to the environment class, and the entry point to the environment configuration class.

Note

The gymnasium registry is a global registry. Hence, it is important to ensure that the environment names are unique. Otherwise, the registry will throw an error when registering the environment.

Manager-Based Environments#

For manager-based environments, the following shows the registration call for the cartpole environment in the isaaclab_tasks.core.cartpole sub-package:

from . import agents

        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerCfg",
        "torchrl_cfg_entry_point": f"{agents.__name__}.torchrl_ppo_cfg:CartpolePPOCfg",
        "default_agent": "rsl_rl",
        "rsl_rl_with_symmetry_cfg_entry_point": (
            f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerWithSymmetryCfg"
        ),
        "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml",
        "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml",
    },
)

gym.register(
    id="Isaac-Cartpole-Camera",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    disable_env_checker=True,

The id argument is the name of the environment. As a convention, we name all the environments with the prefix Isaac- to make it easier to search for them in the registry. The name of the environment is typically followed by the name of the task, and then the name of the robot. For instance, for legged locomotion with ANYmal C on flat terrain, the environment is called IsaacContrib-Velocity-Flat-AnymalC. The version number v<N> is typically used to specify different variations of the same environment. Otherwise, the names of the environments can become too long and difficult to read.

The entry_point argument is the entry point to the environment class. The entry point is a string of the form <module>:<class>. In the case of the cartpole environment, the entry point is isaaclab.envs:ManagerBasedRLEnv. The entry point is used to import the environment class when creating the environment instance.

The env_cfg_entry_point argument specifies the default configuration for the environment. The default configuration is loaded using the isaaclab_tasks.utils.parse_env_cfg() function. It is then passed to the gymnasium.make() function to create the environment instance. The configuration entry point can be both a YAML file or a python configuration class.

Direct Environments#

For direct-based environments, the environment registration follows a similar pattern. Instead of registering the environment’s entry point as the ManagerBasedRLEnv class, we register the environment’s entry point as the implementation class of the environment. Additionally, we add the suffix -Direct to the environment name to differentiate it from the manager-based environments.

As an example, the following shows the registration call for the cartpole environment in the isaaclab_tasks.core.cartpole sub-package:

from . import agents

        "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml",
    },
)

gym.register(
    id="Isaac-Cartpole-Camera-Direct",
    entry_point=f"{__name__}.cartpole_direct_camera_env:CartpoleCameraEnv",
    disable_env_checker=True,
    kwargs={
        "env_cfg_entry_point": f"{__name__}.cartpole_direct_camera_env_cfg:CartpoleCameraEnvCfg",
        "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml",
        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpoleCameraDirectPPORunnerCfg",

Creating the environment#

To inform the gym registry with all the environments provided by the isaaclab_tasks extension, we must import the module at the start of the script. This will execute the __init__.py file which iterates over all the sub-packages and registers their respective environments.

import isaaclab_tasks  # noqa: F401

In this tutorial, the task name is read from the command line. The task name is used to parse the default configuration as well as to create the environment instance. In addition, other parsed command line arguments such as the number of environments, the simulation device, and whether to render, are used to override the default configuration.

    env_cfg, _ = resolve_task_config(args_cli.task, "")
    apply_env_overrides(args_cli, env_cfg)
    pre_launch_video_config(env_cfg, args_cli)
    # reject unsupported configurations before launching Kit or initializing a native physics backend
    try:
        validate(env_cfg)
    except (TypeError, ValueError) as exc:
        raise SystemExit(f"Invalid environment configuration: {exc}") from None

    with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup:
        log_dir = os.path.abspath(os.path.join("logs", f"{policy}_agent", normalize_task_name(args_cli.task)))
        apply_video_recording(env_cfg, log_dir, args_cli, subdir="play")
        env = gym.make(args_cli.task, cfg=env_cfg)

Once creating the environment, the rest of the execution follows the standard resetting and stepping.

The Code Execution#

Now that we have gone through the code, let’s run the script and see the result:

uv (Recommended)

uv run python scripts/environments/random_agent.py --task Isaac-Cartpole --num_envs 32 --viz kit

This should open a stage with everything similar to the Creating a Manager-Based RL Environment tutorial. To stop the simulation, you can either close the window, or press Ctrl+C in the terminal.

In addition, you can also change the simulation device from GPU to CPU by setting the value of the --device flag explicitly: