Playground environments¶

Overview¶

The repository https://github.com/google-deepmind/mujoco_playground provides the MuJoCo Playground environments.

These environments can be easily loaded and configured by calling a single function provided with this library. Such function also makes it possible to configure the environment from the command line arguments or from its parameters.

Important

MuJoCo Playground 1.0.0 or later is required (pip install "playground>=1.0.0", Python 3.12 or later). Since version 1.0.0, MuJoCo Playground environments are implemented on top of PyTorch and MuJoCo Warp (the previous JAX/MJX-based environments are not supported). Therefore, PyTorch is required regardless of the ML framework used by skrl.

Note

The command line arguments has priority over the function parameters.

Run the following command to list all the available MuJoCo Playground environments:

python -c "from mujoco_playground.rl import env_registry; print(list(env_registry._REGISTRY))"

Usage¶

The following snippets show how to load MuJoCo Playground environments:

# import the environment loader
from skrl.envs.loaders.torch import load_playground_env

# load environment
env = load_playground_env(task_name="CartpoleBalance", num_envs=1024, cfg_overrides={"episode_length_secs": 3.0})

# import the environment loader
from skrl.envs.loaders.jax import load_playground_env

# load environment
env = load_playground_env(task_name="CartpoleBalance", num_envs=1024, cfg_overrides={"episode_length_secs": 3.0})

# import the environment loader
from skrl.envs.loaders.warp import load_playground_env

# load environment
env = load_playground_env(task_name="CartpoleBalance", num_envs=1024, cfg_overrides={"episode_length_secs": 3.0})

# import the environment loader
from skrl.envs.loaders.torch import load_playground_env

# load environment
env = load_playground_env()

# import the environment loader
from skrl.envs.loaders.jax import load_playground_env

# load environment
env = load_playground_env()

# import the environment loader
from skrl.envs.loaders.warp import load_playground_env

# load environment
env = load_playground_env()

Run the main script passing the configuration as command line arguments. For example:

python script.py --task CartpoleBalance --num_envs 1024 --device cuda:0

API¶

PyTorch¶

skrl.envs.loaders.torch.load_playground_env(*, task_name: str = '', num_envs: int | None = None, device: str | None = None, cfg_overrides: dict[str, Any] | None = None, show_cfg: bool = True, parser: ArgumentParser | None = None)[source]¶

Load a MuJoCo Playground environment.

MuJoCo Playground: https://playground.mujoco.org

Note

MuJoCo Playground 1.0.0 or later (implemented on top of PyTorch and MuJoCo Warp) is required.

This function includes the definition and parsing of the following command line arguments:

--task: Name of the task.
--num_envs: Number of environments to simulate.
--device: Device (e.g.: "cpu", "cuda:0") used to simulate the environments.
--seed: Random seed.

Parameters:

task_name – The name of the task. If not specified, the task name is taken from the command line argument (--task TASK_NAME). Command line argument has priority over function parameter if both are specified.
num_envs – Number of parallel environments to create. If not specified, the number of environments is taken from the command line argument (--num_envs N). Command line argument has priority over function parameter if both are specified.
device – Device used to simulate the environments. If neither the function parameter nor the command line argument (--device DEVICE) is specified, skrl’s default PyTorch device (skrl.config.torch.device) is used. Command line argument has priority over function parameter if both are specified.
cfg_overrides – Task configuration overrides, where the keys are (dotted) attribute paths of the task configuration (e.g.: {"episode_length_secs": 5.0, "sim.impl": "cpp"}). The "sim.nworld" and "sim.xpu_device" keys are not allowed (use num_envs and device instead).
show_cfg – Whether to print the configuration.
parser – The argument parser to use. If not specified, a new argument parser will be created.

Returns:

MuJoCo Playground environment.

Raises:

ValueError – The task name has not been defined, neither by the function parameter nor by the command line arguments, or the task name is invalid.
ValueError – The number of environments has not been defined, neither by the function parameter nor by the command line arguments.
ValueError – A configuration override sets the number of environments or the device.
AttributeError – A configuration override does not exist in the task configuration.

JAX¶

skrl.envs.loaders.jax.load_playground_env(*, task_name: str = '', num_envs: int | None = None, device: str | None = None, cfg_overrides: dict[str, Any] | None = None, show_cfg: bool = True, parser: ArgumentParser | None = None)[source]¶

Load a MuJoCo Playground environment.

MuJoCo Playground: https://playground.mujoco.org

Note

MuJoCo Playground 1.0.0 or later (implemented on top of PyTorch and MuJoCo Warp) is required.

This function includes the definition and parsing of the following command line arguments:

--task: Name of the task.
--num_envs: Number of environments to simulate.
--device: Device (e.g.: "cpu", "cuda:0") used to simulate the environments.
--seed: Random seed.

Parameters:

task_name – The name of the task. If not specified, the task name is taken from the command line argument (--task TASK_NAME). Command line argument has priority over function parameter if both are specified.
num_envs – Number of parallel environments to create. If not specified, the number of environments is taken from the command line argument (--num_envs N). Command line argument has priority over function parameter if both are specified.
device – Device used to simulate the environments. If neither the function parameter nor the command line argument (--device DEVICE) is specified, skrl’s default JAX device (skrl.config.jax.device) is used. Command line argument has priority over function parameter if both are specified.
cfg_overrides – Task configuration overrides, where the keys are (dotted) attribute paths of the task configuration (e.g.: {"episode_length_secs": 5.0, "sim.impl": "cpp"}). The "sim.nworld" and "sim.xpu_device" keys are not allowed (use num_envs and device instead).
show_cfg – Whether to print the configuration.
parser – The argument parser to use. If not specified, a new argument parser will be created.

Returns:

MuJoCo Playground environment.

Raises:

ValueError – The task name has not been defined, neither by the function parameter nor by the command line arguments, or the task name is invalid.
ValueError – The number of environments has not been defined, neither by the function parameter nor by the command line arguments.
ValueError – A configuration override sets the number of environments or the device.
AttributeError – A configuration override does not exist in the task configuration.

Warp¶

skrl.envs.loaders.warp.load_playground_env(*, task_name: str = '', num_envs: int | None = None, device: str | None = None, cfg_overrides: dict[str, Any] | None = None, show_cfg: bool = True, parser: ArgumentParser | None = None)[source]¶

Load a MuJoCo Playground environment.

MuJoCo Playground: https://playground.mujoco.org

Note

MuJoCo Playground 1.0.0 or later (implemented on top of PyTorch and MuJoCo Warp) is required.

This function includes the definition and parsing of the following command line arguments:

--task: Name of the task.
--num_envs: Number of environments to simulate.
--device: Device (e.g.: "cpu", "cuda:0") used to simulate the environments.
--seed: Random seed.

Parameters:

task_name – The name of the task. If not specified, the task name is taken from the command line argument (--task TASK_NAME). Command line argument has priority over function parameter if both are specified.
num_envs – Number of parallel environments to create. If not specified, the number of environments is taken from the command line argument (--num_envs N). Command line argument has priority over function parameter if both are specified.
device – Device used to simulate the environments. If neither the function parameter nor the command line argument (--device DEVICE) is specified, skrl’s default Warp device (skrl.config.warp.device) is used. Command line argument has priority over function parameter if both are specified.
cfg_overrides – Task configuration overrides, where the keys are (dotted) attribute paths of the task configuration (e.g.: {"episode_length_secs": 5.0, "sim.impl": "cpp"}). The "sim.nworld" and "sim.xpu_device" keys are not allowed (use num_envs and device instead).
show_cfg – Whether to print the configuration.
parser – The argument parser to use. If not specified, a new argument parser will be created.

Returns:

MuJoCo Playground environment.

Raises:

ValueError – The task name has not been defined, neither by the function parameter nor by the command line arguments, or the task name is invalid.
ValueError – The number of environments has not been defined, neither by the function parameter nor by the command line arguments.
ValueError – A configuration override sets the number of environments or the device.
AttributeError – A configuration override does not exist in the task configuration.