Wrapping (single-agent)



This library works with a common API to interact with the following RL environments:

To operate with them and to support interoperability between these non-compatible interfaces, a wrapping mechanism is provided as shown in the diagram below


Environment wrapping Environment wrapping

Usage

# import the environment wrapper and loader
from skrl.envs.wrappers.torch import wrap_env
from skrl.envs.loaders.torch import load_isaaclab_env

# load the environment
env = load_isaaclab_env(task_name="Isaac-Cartpole-Direct-v0")

# wrap the environment
env = wrap_env(env)  # or 'env = wrap_env(env, wrapper="isaaclab")'

API (PyTorch)

skrl.envs.wrappers.torch.wrap_env(env: Any, wrapper: str = 'auto', verbose: bool = True) Wrapper | MultiAgentEnvWrapper

Wrap an environment to use a common interface

Example:

>>> from skrl.envs.wrappers.torch import wrap_env
>>>
>>> # assuming that there is an environment called "env"
>>> env = wrap_env(env)
Parameters:
  • env (gym.Env, gymnasium.Env, dm_env.Environment or VecTask) – The environment to be wrapped

  • wrapper (str, optional) –

    The type of wrapper to use (default: "auto"). If "auto", the wrapper will be automatically selected based on the environment class. The supported wrappers are described in the following table:

    Single-agent environments

    Environment

    Wrapper tag

    OpenAI Gym

    "gym"

    Gymnasium

    "gymnasium"

    DeepMind

    "dm"

    Brax

    "brax"

    Isaac Lab

    "isaaclab" ("isaaclab-single-agent")

    Isaac Gym preview 2

    "isaacgym-preview2"

    Isaac Gym preview 3

    "isaacgym-preview3"

    Isaac Gym preview 4

    "isaacgym-preview4"

    Omniverse Isaac Gym

    "omniverse-isaacgym"

    Robosuite

    "robosuite"

    Multi-agent environments

    Environment

    Wrapper tag

    Petting Zoo

    "pettingzoo"

    Isaac Lab

    "isaaclab" ("isaaclab-multi-agent")

    Bi-DexHands

    "bidexhands"

  • verbose (bool, optional) – Whether to print the wrapper type (default: True)

Raises:

ValueError – Unknown wrapper type

Returns:

Wrapped environment

Return type:

Wrapper or MultiAgentEnvWrapper


API (JAX)

skrl.envs.wrappers.jax.wrap_env(env: Any, wrapper: str = 'auto', verbose: bool = True) Wrapper | MultiAgentEnvWrapper

Wrap an environment to use a common interface

Example:

>>> from skrl.envs.wrappers.jax import wrap_env
>>>
>>> # assuming that there is an environment called "env"
>>> env = wrap_env(env)
Parameters:
  • env (gym.Env, gymnasium.Env, dm_env.Environment or VecTask) – The environment to be wrapped

  • wrapper (str, optional) –

    The type of wrapper to use (default: "auto"). If "auto", the wrapper will be automatically selected based on the environment class. The supported wrappers are described in the following table:

    Single-agent environments

    Environment

    Wrapper tag

    OpenAI Gym

    "gym"

    Gymnasium

    "gymnasium"

    Brax

    "brax"

    Isaac Lab

    "isaaclab" ("isaaclab-single-agent")

    Isaac Gym preview 2

    "isaacgym-preview2"

    Isaac Gym preview 3

    "isaacgym-preview3"

    Isaac Gym preview 4

    "isaacgym-preview4"

    Omniverse Isaac Gym

    "omniverse-isaacgym"

    Multi-agent environments

    Environment

    Wrapper tag

    Petting Zoo

    "pettingzoo"

    Isaac Lab

    "isaaclab" ("isaaclab-multi-agent")

    Bi-DexHands

    "bidexhands"

  • verbose (bool, optional) – Whether to print the wrapper type (default: True)

Raises:

ValueError – Unknown wrapper type

Returns:

Wrapped environment

Return type:

Wrapper or MultiAgentEnvWrapper


Internal API (PyTorch)

class skrl.envs.wrappers.torch.Wrapper(env: Any)

Bases: object

__init__(env: Any) None

Base wrapper class for RL environments

Parameters:

env (Any supported RL environment) – The environment to wrap

property action_space: gym.Space

Action space

close() None

Close the environment

Raises:

NotImplementedError – Not implemented

property device: torch.device

The device used by the environment

If the wrapped environment does not have the device property, the value of this property will be "cuda" or "cpu" depending on the device availability

property num_agents: int

Number of agents

If the wrapped environment does not have the num_agents property, it will be set to 1

property num_envs: int

Number of environments

If the wrapped environment does not have the num_envs property, it will be set to 1

property observation_space: gym.Space

Observation space

render(*args, **kwargs) Any

Render the environment

Raises:

NotImplementedError – Not implemented

Returns:

Any value from the wrapped environment

Return type:

any

reset() Tuple[torch.Tensor, Any]

Reset the environment

Raises:

NotImplementedError – Not implemented

Returns:

Observation, info

Return type:

torch.Tensor and any other info

state() torch.Tensor

Get the environment state

Raises:

NotImplementedError – Not implemented

Returns:

State

Return type:

torch.Tensor

property state_space: gym.Space | None

State space

If the wrapped environment does not have the state_space property, None will be returned

step(actions: torch.Tensor) Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Any]

Perform a step in the environment

Parameters:

actions (torch.Tensor) – The actions to perform

Raises:

NotImplementedError – Not implemented

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of torch.Tensor and any other info

class skrl.envs.wrappers.torch.IsaacLabWrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Isaac Lab environment wrapper

Parameters:

env (Any supported Isaac Lab environment) – The environment to wrap

property action_space: gymnasium.Space

Action space

close() None

Close the environment

property observation_space: gymnasium.Space

Observation space

render(*args, **kwargs) None

Render the environment

reset() Tuple[torch.Tensor, Any]

Reset the environment

Returns:

Observation, info

Return type:

torch.Tensor and any other info

property state_space: gymnasium.Space | None

State space

step(actions: torch.Tensor) Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Any]

Perform a step in the environment

Parameters:

actions (torch.Tensor) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of torch.Tensor and any other info

class skrl.envs.wrappers.torch.OmniverseIsaacGymWrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Omniverse Isaac Gym environment wrapper

Parameters:

env (Any supported Omniverse Isaac Gym environment) – The environment to wrap

close() None

Close the environment

render(*args, **kwargs) None

Render the environment

reset() Tuple[torch.Tensor, Any]

Reset the environment

Returns:

Observation, info

Return type:

torch.Tensor and any other info

run(trainer: omni.isaac.gym.vec_env.vec_env_mt.TrainerMT | None = None) None

Run the simulation in the main thread

This method is valid only for the Omniverse Isaac Gym multi-threaded environments

Parameters:

trainer (omni.isaac.gym.vec_env.vec_env_mt.TrainerMT, optional) – Trainer which should implement a run method that initiates the RL loop on a new thread

step(actions: torch.Tensor) Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Any]

Perform a step in the environment

Parameters:

actions (torch.Tensor) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of torch.Tensor and any other info

class skrl.envs.wrappers.torch.IsaacGymPreview3Wrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Isaac Gym environment (preview 3) wrapper

Parameters:

env (Any supported Isaac Gym environment (preview 3) environment) – The environment to wrap

close() None

Close the environment

render(*args, **kwargs) None

Render the environment

reset() Tuple[torch.Tensor, Any]

Reset the environment

Returns:

Observation, info

Return type:

torch.Tensor and any other info

property state_space: gym.Space | None

State space

step(actions: torch.Tensor) Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Any]

Perform a step in the environment

Parameters:

actions (torch.Tensor) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of torch.Tensor and any other info

class skrl.envs.wrappers.torch.IsaacGymPreview2Wrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Isaac Gym environment (preview 2) wrapper

Parameters:

env (Any supported Isaac Gym environment (preview 2) environment) – The environment to wrap

close() None

Close the environment

render(*args, **kwargs) None

Render the environment

reset() Tuple[torch.Tensor, Any]

Reset the environment

Returns:

Observation, info

Return type:

torch.Tensor and any other info

step(actions: torch.Tensor) Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Any]

Perform a step in the environment

Parameters:

actions (torch.Tensor) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of torch.Tensor and any other info

class skrl.envs.wrappers.torch.GymWrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

OpenAI Gym environment wrapper

Parameters:

env (Any supported OpenAI Gym environment) – The environment to wrap

property action_space: gym.Space

Action space

close() None

Close the environment

property observation_space: gym.Space

Observation space

render(*args, **kwargs) Any

Render the environment

reset() Tuple[torch.Tensor, Any]

Reset the environment

Returns:

Observation, info

Return type:

torch.Tensor and any other info

step(actions: torch.Tensor) Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Any]

Perform a step in the environment

Parameters:

actions (torch.Tensor) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of torch.Tensor and any other info

class skrl.envs.wrappers.torch.GymnasiumWrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Gymnasium environment wrapper

Parameters:

env (Any supported Gymnasium environment) – The environment to wrap

property action_space: gymnasium.Space

Action space

close() None

Close the environment

property observation_space: gymnasium.Space

Observation space

render(*args, **kwargs) Any

Render the environment

reset() Tuple[torch.Tensor, Any]

Reset the environment

Returns:

Observation, info

Return type:

torch.Tensor and any other info

step(actions: torch.Tensor) Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Any]

Perform a step in the environment

Parameters:

actions (torch.Tensor) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of torch.Tensor and any other info

class skrl.envs.wrappers.torch.DeepMindWrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

DeepMind environment wrapper

Parameters:

env (Any supported DeepMind environment) – The environment to wrap

_observation_to_tensor(observation: Any, spec: Any | None = None) torch.Tensor

Convert the DeepMind observation to a flat tensor

Parameters:

observation (Any supported DeepMind observation) – The DeepMind observation to convert to a tensor

Raises:

ValueError if the observation spec type is not supported

Returns:

The observation as a flat tensor

Return type:

torch.Tensor

_spec_to_space(spec: Any) gym.Space

Convert the DeepMind spec to a Gym space

Parameters:

spec (Any supported DeepMind spec) – The DeepMind spec to convert

Raises:

ValueError if the spec type is not supported

Returns:

The Gym space

Return type:

gym.Space

_tensor_to_action(actions: torch.Tensor) Any

Convert the action to the DeepMind expected format

Parameters:

actions (torch.Tensor) – The actions to perform

Raises:

ValueError – If the action space type is not supported

Returns:

The action in the DeepMind expected format

Return type:

Any supported DeepMind action

property action_space: gym.Space

Action space

close() None

Close the environment

property observation_space: gym.Space

Observation space

render(*args, **kwargs) ndarray

Render the environment

OpenCV is used to render the environment. Install OpenCV with pip install opencv-python

reset() Tuple[torch.Tensor, Any]

Reset the environment

Returns:

The state of the environment

Return type:

torch.Tensor

step(actions: torch.Tensor) Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Any]

Perform a step in the environment

Parameters:

actions (torch.Tensor) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of torch.Tensor and any other info

class skrl.envs.wrappers.torch.BraxWrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Brax environment wrapper

Parameters:

env (Any supported Brax environment) – The environment to wrap

property action_space: gymnasium.Space

Action space

close() None

Close the environment

property observation_space: gymnasium.Space

Observation space

render(*args, **kwargs) None

Render the environment

reset() Tuple[torch.Tensor, Any]

Reset the environment

Returns:

Observation, info

Return type:

torch.Tensor and any other info

step(actions: torch.Tensor) Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Any]

Perform a step in the environment

Parameters:

actions (torch.Tensor) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of torch.Tensor and any other info

class skrl.envs.wrappers.torch.RobosuiteWrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Robosuite environment wrapper

Parameters:

env (Any supported robosuite environment) – The environment to wrap

_observation_to_tensor(observation: Any, spec: Any | None = None) torch.Tensor

Convert the observation to a flat tensor

Parameters:

observation (Any supported observation) – The observation to convert to a tensor

Raises:

ValueError if the observation spec type is not supported

Returns:

The observation as a flat tensor

Return type:

torch.Tensor

_spec_to_space(spec: Any) gym.Space

Convert the robosuite spec to a Gym space

Parameters:

spec (Any supported robosuite spec) – The robosuite spec to convert

Raises:

ValueError if the spec type is not supported

Returns:

The Gym space

Return type:

gym.Space

_tensor_to_action(actions: torch.Tensor) Any

Convert the action to the robosuite expected format

Parameters:

actions (torch.Tensor) – The actions to perform

Raises:

ValueError – If the action space type is not supported

Returns:

The action in the robosuite expected format

Return type:

Any supported robosuite action

property action_space: gym.Space

Action space

close() None

Close the environment

property observation_space: gym.Space

Observation space

render(*args, **kwargs) None

Render the environment

reset() Tuple[torch.Tensor, Any]

Reset the environment

Returns:

The state of the environment

Return type:

torch.Tensor

property state_space: gym.Space

State space

An alias for the observation_space property

step(actions: torch.Tensor) Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Any]

Perform a step in the environment

Parameters:

actions (torch.Tensor) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of torch.Tensor and any other info


Internal API (JAX)

class skrl.envs.wrappers.jax.Wrapper(env: Any)

Bases: object

__init__(env: Any) None

Base wrapper class for RL environments

Parameters:

env (Any supported RL environment) – The environment to wrap

property action_space: gym.Space

Action space

close() None

Close the environment

Raises:

NotImplementedError – Not implemented

property device: jax.Device

The device used by the environment

If the wrapped environment does not have the device property, the value of this property will be "cuda" or "cpu" depending on the device availability

property num_agents: int

Number of agents

If the wrapped environment does not have the num_agents property, it will be set to 1

property num_envs: int

Number of environments

If the wrapped environment does not have the num_envs property, it will be set to 1

property observation_space: gym.Space

Observation space

render(*args, **kwargs) Any

Render the environment

Raises:

NotImplementedError – Not implemented

Returns:

Any value from the wrapped environment

Return type:

any

reset() Tuple[ndarray | jax.Array, Any]

Reset the environment

Raises:

NotImplementedError – Not implemented

Returns:

Observation, info

Return type:

np.ndarray or jax.Array and any other info

state() ndarray | jax.Array

Get the environment state

Raises:

NotImplementedError – Not implemented

Returns:

State

Return type:

np.ndarray or jax.Array

property state_space: gym.Space | None

State space

If the wrapped environment does not have the state_space property, None will be returned

step(actions: ndarray | jax.Array) Tuple[ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, Any]

Perform a step in the environment

Parameters:

actions (np.ndarray or jax.Array) – The actions to perform

Raises:

NotImplementedError – Not implemented

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of np.ndarray or jax.Array and any other info

class skrl.envs.wrappers.jax.IsaacLabWrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Isaac Lab environment wrapper

Parameters:

env (Any supported Isaac Lab environment) – The environment to wrap

property action_space: gymnasium.Space

Action space

close() None

Close the environment

property observation_space: gymnasium.Space

Observation space

render(*args, **kwargs) None

Render the environment

reset() Tuple[ndarray | jax.Array, Any]

Reset the environment

Returns:

Observation, info

Return type:

np.ndarray or jax.Array and any other info

property state_space: gymnasium.Space | None

State space

step(actions: ndarray | jax.Array) Tuple[ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, Any]

Perform a step in the environment

Parameters:

actions (np.ndarray or jax.Array) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of np.ndarray or jax.Array and any other info

class skrl.envs.wrappers.jax.OmniverseIsaacGymWrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Omniverse Isaac Gym environment wrapper

Parameters:

env (Any supported Omniverse Isaac Gym environment) – The environment to wrap

close() None

Close the environment

render(*args, **kwargs) None

Render the environment

reset() Tuple[ndarray | jax.Array, Any]

Reset the environment

Returns:

Observation, info

Return type:

np.ndarray or jax.Array and any other info

run(trainer: omni.isaac.gym.vec_env.vec_env_mt.TrainerMT | None = None) None

Run the simulation in the main thread

This method is valid only for the Omniverse Isaac Gym multi-threaded environments

Parameters:

trainer (omni.isaac.gym.vec_env.vec_env_mt.TrainerMT, optional) – Trainer which should implement a run method that initiates the RL loop on a new thread

step(actions: ndarray | jax.Array) Tuple[ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, Any]

Perform a step in the environment

Parameters:

actions (np.ndarray or jax.Array) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of np.ndarray or jax.Array and any other info

class skrl.envs.wrappers.jax.IsaacGymPreview3Wrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Isaac Gym environment (preview 3) wrapper

Parameters:

env (Any supported Isaac Gym environment (preview 3) environment) – The environment to wrap

close() None

Close the environment

render(*args, **kwargs) None

Render the environment

reset() Tuple[ndarray | jax.Array, Any]

Reset the environment

Returns:

Observation, info

Return type:

np.ndarray or jax.Array and any other info

property state_space: gym.Space | None

State space

step(actions: ndarray | jax.Array) Tuple[ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, Any]

Perform a step in the environment

Parameters:

actions (np.ndarray or jax.Array) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of np.ndarray or jax.Array and any other info

class skrl.envs.wrappers.jax.IsaacGymPreview2Wrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Isaac Gym environment (preview 2) wrapper

Parameters:

env (Any supported Isaac Gym environment (preview 2) environment) – The environment to wrap

close() None

Close the environment

render(*args, **kwargs) None

Render the environment

reset() Tuple[ndarray | jax.Array, Any]

Reset the environment

Returns:

Observation, info

Return type:

np.ndarray or jax.Array and any other info

step(actions: ndarray | jax.Array) Tuple[ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, Any]

Perform a step in the environment

Parameters:

actions (np.ndarray or jax.Array) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of np.ndarray or jax.Array and any other info

class skrl.envs.wrappers.jax.GymnasiumWrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Gymnasium environment wrapper

Parameters:

env (Any supported Gymnasium environment) – The environment to wrap

property action_space: gymnasium.Space

Action space

close() None

Close the environment

property observation_space: gymnasium.Space

Observation space

render(*args, **kwargs) Any

Render the environment

reset() Tuple[ndarray | jax.Array, Any]

Reset the environment

Returns:

Observation, info

Return type:

np.ndarray or jax.Array and any other info

step(actions: ndarray | jax.Array) Tuple[ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, Any]

Perform a step in the environment

Parameters:

actions (np.ndarray or jax.Array) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of np.ndarray or jax.Array and any other info

class skrl.envs.wrappers.jax.BraxWrapper(env: Any)

Bases: Wrapper

__init__(env: Any) None

Brax environment wrapper

Parameters:

env (Any supported Brax environment) – The environment to wrap

property action_space: gymnasium.Space

Action space

close() None

Close the environment

property observation_space: gymnasium.Space

Observation space

render(*args, **kwargs) None

Render the environment

reset() Tuple[ndarray | jax.Array, Any]

Reset the environment

Returns:

Observation, info

Return type:

np.ndarray or jax.Array and any other info

step(actions: ndarray | jax.Array) Tuple[ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, ndarray | jax.Array, Any]

Perform a step in the environment

Parameters:

actions (np.ndarray or jax.Array) – The actions to perform

Returns:

Observation, reward, terminated, truncated, info

Return type:

tuple of np.ndarray or jax.Array and any other info