### Setup Checkpoint Directory and Config

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/gait_tracking.ipynb

Creates a checkpoint directory and saves the environment configuration to a JSON file. This is typically done before training starts.

```python
ckpt_path = epath.Path("checkpoints").resolve() / exp_name
ckpt_path.mkdir(parents=True, exist_ok=True)
print(f"Checkpoint path: {ckpt_path}")

with open(ckpt_path / "config.json", "w") as fp:
  json.dump(env_cfg.to_dict(), fp, indent=4)
```

--------------------------------

### Install Pre-requisites

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/vision.ipynb

Installs necessary libraries including warp-lang, mujoco, mujoco_mjx, and brax. Ensure you are in a Colab environment.

```python
#@title Install pre-requisites
!pip install warp-lang
!pip install mujoco
!pip install mujoco_mjx
!pip install brax
```

--------------------------------

### Setup Checkpointing and Training Parameters

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/open_cabinet.ipynb

Initializes the checkpoint path, saves the environment configuration, and defines the training function with PPO parameters.

```python
ckpt_path = epath.Path("checkpoints").resolve() / exp_name
ckpt_path.mkdir(parents=True, exist_ok=True)
print(f"Checkpoint path: {ckpt_path}")

with open(ckpt_path / "config.json", "w") as fp:
  json.dump(env_cfg.to_dict(), fp, indent=4)


def policy_params_fn(current_step, make_policy, params):
  orbax_checkpointer = ocp.PyTreeCheckpointer()
  save_args = orbax_utils.save_args_from_target(params)
  path = ckpt_path / f"{current_step}"
  orbax_checkpointer.save(path, params, force=True, save_args=save_args)

from mujoco_playground.config import manipulation_params

_train_params = manipulation_params.brax_ppo_config(env_name)
train_params = dict(_train_params)
train_params['seed'] = 1
del train_params["network_factory"]

train_fn = functools.partial(
  ppo.train,
  **dict(train_params),
  network_factory=functools.partial(
    ppo_networks.make_ppo_networks,
    policy_hidden_layer_sizes=_train_params.network_factory.policy_hidden_layer_sizes
  ))

```

--------------------------------

### Clone and Install MuJoCo Playground from Source

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/README.md

Clone the repository and install the MuJoCo Playground from source using uv for faster installation. Requires Python 3.10 or later.

```sh
git clone git@github.com:google-deepmind/mujoco_playground.git && cd mujoco_playground
```

```sh
uv venv --python 3.12
```

```sh
source .venv/bin/activate
```

```sh
uv pip install -U "jax[cuda12]" --index-url https://pypi.org/simple
```

```sh
uv --no-config sync --all-extras
```

--------------------------------

### Install Pre-requisites

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/manipulation.ipynb

Installs the necessary MuJoCo, MuJoCo MJX, and Brax libraries. Ensure you are in a Colab environment with GPU acceleration enabled.

```python
#@title Install pre-requisites
!pip install mujoco
!pip install mujoco_mjx
!pip install brax
```

--------------------------------

### Configure and Initialize PPO Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/apollo_joystick.ipynb

Prepares training parameters, defines the training function using `functools.partial`, loads the environment, and starts the training process.

```python
training_params = dict(ppo_params)
del training_params["network_factory"]

train_fn = functools.partial(
  ppo.train,
  **training_params,
  network_factory=functools.partial(
    ppo_networks.make_ppo_networks, **ppo_params.network_factory
  ),
  restore_checkpoint_path=restore_checkpoint_path,
  progress_fn=progress,
  wrap_env_fn=wrapper.wrap_for_brax_training,
  policy_params_fn=policy_params_fn,
  randomization_fn=randomizer,
)

env = registry.load(env_name, config=env_cfg)
eval_env = registry.load(env_name, config=env_cfg)
make_inference_fn, params, _ = train_fn(environment=env, eval_env=eval_env)
if len(plotter.times) > 1:
  print(f"time to jit: {plotter.times[1] - plotter.times[0]}")
  print(f"time to train: {plotter.times[-1] - plotter.times[1]}")
```

--------------------------------

### Initialize Environment and Start Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/open_cabinet.ipynb

Loads the MuJoCo environment, wraps it for training, and initiates the PPO training process with the defined progress callback.

```python
env = manipulation.load(env_name, config=env_cfg)
make_inference_fn, params, _ = train_fn(
    wrap_env=False,
    environment=wrapper.wrap_for_brax_training(env,
                                               episode_length=env_cfg.episode_length,
                                               action_repeat=env_cfg.action_repeat),
                                               progress_fn=progress
)
print(f"time to jit: {times[1] - times[0]}")
print(f"time to train: {times[-1] - times[1]}")
```

--------------------------------

### Initialize and Run PPO Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/humanoid_joystick.ipynb

Sets up the PPO training function with specified parameters, environment wrappers, randomization, and checkpointing callbacks. It then loads the environment and starts the training process.

```python
training_params = dict(ppo_params)
del training_params["network_factory"]

train_fn = functools.partial(
  ppo.train,
  **training_params,
  network_factory=functools.partial(
      ppo_networks.make_ppo_networks,
      **ppo_params.network_factory
  ),
  restore_checkpoint_path=restore_checkpoint_path,
  progress_fn=progress,
  wrap_env_fn=wrapper.wrap_for_brax_training,
  policy_params_fn=policy_params_fn,
  randomization_fn=randomizer,
)

env = locomotion.load(env_name, config=env_cfg)
eval_env = locomotion.load(env_name, config=env_cfg)
make_inference_fn, params, _ = train_fn(environment=env, eval_env=eval_env)
if len(times) > 1:
  print(f"time to jit: {times[1] - times[0]}")
  print(f"time to train: {times[-1] - times[1]}")
```

--------------------------------

### Install MuJoCo Playground from PyPI

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/README.md

Install the MuJoCo Playground library directly from the Python Package Index.

```sh
pip install playground
```

--------------------------------

### Initialize and Run PPO Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/getup.ipynb

Sets up the PPO training function with specified parameters, environment wrappers, randomization, and checkpointing callbacks. It then initializes the environment and starts the training process.

```python
training_params = dict(ppo_params)
del training_params["network_factory"]

train_fn = functools.partial(
  ppo.train,
  **training_params,
  network_factory=functools.partial(
      ppo_networks.make_ppo_networks,
      **ppo_params.network_factory
  ),
  restore_checkpoint_path=restore_checkpoint_path,
  progress_fn=progress,
  wrap_env_fn=wrapper.wrap_for_brax_training,
  policy_params_fn=policy_params_fn,
  randomization_fn=randomizer,
)

env = locomotion.load(env_name, config=env_cfg)
eval_env = locomotion.load(env_name, config=env_cfg)
make_inference_fn, params, _ = train_fn(environment=env, eval_env=eval_env)
if len(times) > 1:
  print(f"time to jit: {times[1] - times[0]}")
  print(f"time to train: {times[-1] - times[1]}")

# Make a final plot of reward and success vs WALLCLOCK time.
plt.figure()
plt.ylim([0, 75])
plt.xlabel("wallclock time (s)")
plt.ylabel("reward per episode")
plt.title(f"y={y_data[-1]:.3f}")
plt.errorbar(
    [(t - times[0]).total_seconds() for t in times[:-1]],
    y_data,
    yerr=y_dataerr,
    color="blue",
)
plt.show()
```

--------------------------------

### Install ONNX Runtime and HIDAPI

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/sim2sim/README.md

Install necessary Python packages for running ONNX models and reading joystick input.

```bash
uv pip install onnxruntime hidapi
```

--------------------------------

### Install Development Dependencies and Pre-commit Hooks

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/CONTRIBUTING.md

Install the project's development dependencies and set up pre-commit hooks for code linting and formatting. Run this before submitting a Pull Request.

```shell
pip install -e ".[dev]"
pre-commit install
pre-commit run --all-files
```

--------------------------------

### Start rscope Server

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/README.md

Start the rscope server in a separate terminal to connect with the training visualization.

```bash
python -m rscope
```

--------------------------------

### Import Necessary Libraries

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/joystick_gait_tracking.ipynb

Imports all required libraries for environment setup, training, logging, and visualization.

```python
import functools
import json
from datetime import datetime

import jax
import jax.numpy as jp
import matplotlib.pyplot as plt
import mediapy as media
import mujoco
import numpy as np
import wandb
from brax.training.agents.ppo import networks as ppo_networks
from brax.training.agents.ppo import train as ppo
from etils import epath
from flax.training import orbax_utils
from IPython.display import clear_output, display
from orbax import checkpoint as ocp

from mujoco_playground import BraxEnvWrapper, locomotion
from mujoco_playground._src.gait import draw_joystick_command

# Enable persistent compilation cache.
jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
```

--------------------------------

### Train with PPO (Command Line)

Source: https://context7.com/google-deepmind/mujoco_playground/llms.txt

Command-line examples for training agents using the JAX backend PPO script. Supports various configurations including different backends, domain randomization, and Weights & Biases logging.

```bash
# Basic training with JAX backend
train-jax-ppo --env_name CartpoleBalance --num_timesteps 1000000
```

```bash
# Training with MuJoCo Warp backend (faster on some GPUs)
train-jax-ppo --env_name CartpoleBalance --impl warp
```

```bash
# Training locomotion with domain randomization
train-jax-ppo --env_name Go1JoystickFlatTerrain \
    --domain_randomization \
    --num_timesteps 50000000 \
    --num_envs 4096
```

```bash
# Training with Weights & Biases logging
train-jax-ppo --env_name PandaPickCube \
    --use_wandb \
    --num_timesteps 10000000
```

```bash
# Load checkpoint and continue training
train-jax-ppo --env_name CartpoleBalance \
    --load_checkpoint_path ./logs/CartpoleBalance-20240101/checkpoints
```

```bash
# Play only (no training, just visualize)
train-jax-ppo --env_name CheetahRun \
    --play_only \
    --load_checkpoint_path ./logs/CheetahRun-trained/checkpoints
```

--------------------------------

### Setup Checkpoint Directory and Save Configuration

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/cube_rotate.ipynb

Creates a checkpoint directory based on the experiment name and saves the environment configuration to a JSON file within that directory.

```python
ckpt_path = epath.Path("checkpoints").resolve() / exp_name
ckpt_path.mkdir(parents=True, exist_ok=True)
print(f"{ckpt_path}")

with open(ckpt_path / "config.json", "w") as fp:
  json.dump(env_cfg.to_dict(), fp, indent=4)
```

--------------------------------

### Load Latest Checkpoint and Setup Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/inference.ipynb

Finds the latest checkpoint, configures the PPO network factory, and sets up the training function with specified parameters. Restores from the latest checkpoint.

```python
latest_ckpts = list(CKPT_PATH.glob("*\n"))
if not latest_ckpts:
  raise ValueError("No checkpoints found")
# Remove .json files.
latest_ckpts = [
    ckpt for ckpt in latest_ckpts if not ckpt.name.endswith(".json")
]
latest_ckpts.sort(key=lambda x: int(x.name))
print(f"Found {len(latest_ckpts)} checkpoints")
latest_ckpt = latest_ckpts[-1]
restore_checkpoint_path = latest_ckpt
print(f"Restoring from {restore_checkpoint_path}")

make_networks_factory = functools.partial(
    ppo_networks.make_ppo_networks,
    policy_hidden_layer_sizes=(128, 128, 128, 128),
)

train_fn = functools.partial(
    ppo.train,
    num_timesteps=0,
    episode_length=1000,
    normalize_observations=True,
    restore_checkpoint_path=restore_checkpoint_path,
    network_factory=make_networks_factory,
    num_envs=2,
)
make_inference_fn, params, _ = train_fn(environment=BraxEnvWrapper(env))
```

--------------------------------

### Train with RSL-RL (Command Line)

Source: https://context7.com/google-deepmind/mujoco_playground/llms.txt

Command-line examples for training locomotion tasks using the PyTorch-based RSL-RL. Includes options for Warp implementation and multi-GPU training.

```bash
# Train with RSL-RL
train-rsl-ppo --env_name BerkeleyHumanoidJoystickFlatTerrain \
    --impl warp \
    --num_envs 4096
```

```bash
# Multi-GPU training
torchrun --nproc_per_node=4 train-rsl-ppo \
    --env_name Go1JoystickFlatTerrain \
    --multi_gpu \
    --num_envs 8192
```

--------------------------------

### Setup Checkpoint Directory and Save Configuration

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb

Creates a checkpoint directory for the specified environment and saves the environment configuration to 'config.json'. This is useful for tracking environment settings during training.

```python
ckpt_path = epath.Path("checkpoints").resolve() / env_name
ckpt_path.mkdir(parents=True, exist_ok=True)
print(f"{ckpt_path}")

with open(ckpt_path / "config.json", "w") as fp:
  json.dump(env_cfg.to_dict(), fp, indent=4)
```

--------------------------------

### JIT-Compiled Environment Rollouts with JAX

Source: https://context7.com/google-deepmind/mujoco_playground/llms.txt

For efficient training, JIT-compile the environment's reset and step functions and run vectorized rollouts using `jax.lax.scan`. This example demonstrates how to perform a rollout and calculate the total reward.

```python
import jax
import jax.numpy as jp
from mujoco_playground import registry

env = registry.load("CheetahRun")

# JIT compile for speed
jit_reset = jax.jit(env.reset)
jit_step = jax.jit(env.step)

# Warmup JIT compilation
rng = jax.random.PRNGKey(0)
state = jit_reset(rng)
action = jp.zeros(env.action_size)
_ = jit_step(state, action)

# Run a rollout
def rollout(rng, num_steps=1000):
    state = jit_reset(rng)

    def step_fn(carry, _):
        state, rng = carry
        rng, act_rng = jax.random.split(rng)
        action = jax.random.uniform(act_rng, (env.action_size,), minval=-1, maxval=1)
        next_state = jit_step(state, action)
        return (next_state, rng), next_state.reward

    (final_state, _), rewards = jax.lax.scan(
        step_fn, (state, rng), None, length=num_steps
    )
    return final_state, rewards

final_state, rewards = rollout(rng)
print(f"Total reward: {rewards.sum():.2f}")
print(f"Episode length: {len(rewards)}")
```

--------------------------------

### Run Go1 Joystick Controller

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/sim2sim/README.md

Execute the Python script to start the Go1 joystick controller simulation.

```bash
python play_go1_joystick.py
```

--------------------------------

### Environment Setup and Training Execution

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/cube_reorient.ipynb

Loads the specified MuJoCo environment and its evaluation counterpart, then initiates the training process using the configured `train_fn`. Prints training and JIT compilation times if available.

```python
env = manipulation.load(env_name, config=env_cfg)
eval_env = manipulation.load(env_name, config=env_cfg)
make_inference_fn, params, _ = train_fn(environment=env, eval_env=eval_env)
if len(times) > 1:
  print(f"time to jit: {times[1] - times[0]}")
  print(f"time to train: {times[-1] - times[1]}")
```

--------------------------------

### Configure PPO Training Parameters (Commented Out)

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/cube_reorient.ipynb

Provides commented-out examples for adjusting PPO training parameters such as `num_timesteps`, `num_evals`, observation noise, and learning rate. These can be uncommented and modified for different training runs.

```python
## First run.
# ppo_params.num_timesteps = 200_000_000
# ppo_params.num_evals = 20

## Second run (add torque limits & frictionloss).
# ppo_params.num_timesteps = 100_000_000
# ppo_params.num_evals = 10
# env_cfg.obs_noise.random_ori_injection_prob = 0.1
# ppo_params.learning_rate = 1e-4
```

--------------------------------

### Setup Weights & Biases Logging

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/cube_rotate.ipynb

Initializes Weights & Biases for experiment tracking if USE_WANDB is set to True. It logs the project name and environment configuration.

```python
# Setup wandb logging.
USE_WANDB = False

if USE_WANDB:
  wandb.init(project="mjxrl", config=env_cfg)
  wandb.config.update({
      "env_name": env_name,
  })
```

--------------------------------

### Setup Inference and Environment

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/getup.ipynb

Initializes the inference function and the evaluation environment. JIT compilation is used for performance optimization. The environment configuration is modified to set `drop_from_height_prob` to 1.0.

```python
inference_fn = make_inference_fn(params, deterministic=True)
jit_inference_fn = jax.jit(inference_fn)

env_cfg.drop_from_height_prob = 1.0
eval_env = locomotion.load(env_name, config=env_cfg)
jit_reset = jax.jit(eval_env.reset)
jit_step = jax.jit(eval_env.step)
```

--------------------------------

### Import Necessary Libraries for Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/position_cube.ipynb

Imports core libraries for JAX, training agents (PPO), environment wrappers, and checkpointing. Ensure these libraries are installed.

```python
import functools
import json
from datetime import datetime

import jax
import matplotlib.pyplot as plt
import mediapy as media
from brax.training.agents.ppo import networks as ppo_networks
from brax.training.agents.ppo import train as ppo
from etils import epath
from flax.training import orbax_utils
from IPython.display import clear_output, display
from orbax import checkpoint as ocp

from mujoco_playground import BraxEnvWrapper, manipulation

# Enable persistent compilation cache.
jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
```

--------------------------------

### Commented Environment Configuration Options

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/joystick.ipynb

Provides commented-out examples for modifying environment configuration, such as episode length, perturbation settings, and command parameters.

```python
# env_cfg.episode_length = 500  # Shorten episode length so we don't go out of boudns of the hfield.
# env_cfg.pert_config.velocity_kick = [3.0, 6.0]
# env_cfg.pert_config.kick_wait_times = [5.0, 15.0]
# env_cfg.command_config.a = [1.5, 0.8, 2*jp.pi]
```

--------------------------------

### Simulate Environment with ONNX Policy

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/brax_network_to_onnx.ipynb

A commented-out example demonstrating how to run a simulation within a JAX environment using an ONNX-converted policy. It sets up a command, resets the environment, and then iteratively steps through the environment using actions predicted by the ONNX model.

```python
# Test the policy.

# env_cfg = locomotion.get_default_config(env_name)
# env_cfg.init_from_crouch = 0.0
# env = locomotion.load(env_name, config=env_cfg)
# env_cfg = manipulation.get_default_config(env_name)
# env = manipulation.load(env_name, config=env_cfg)
# jit_reset = jax.jit(env.reset)
# jit_step = jax.jit(env.step)

x = 0.8
y = 0.0
yaw = 0.3
command = jp.array([x, y, yaw])
# actions = []

states = [state := jit_reset(jax.random.PRNGKey(555))]
state.info["command"] = command
for _ in range(env_cfg.episode_length):
  onnx_input = {'obs': np.array(state.obs["state"].reshape(1, -1))}
  action = m.run(output_names, onnx_input)[0][0]
  state = jit_step(state, jp.array(action))
  state.info["command"] = command
  states.append(state)
  # actions.append(state.info["motor_targets"])
  # actions.append(action)
  if state.done:
    print("Unexpected termination.")
    break
```

--------------------------------

### Execute PPO Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/vision.ipynb

This code block initializes the evaluation environment and then calls the `train_fn` to start the PPO training process. It prints the time taken for JIT compilation and the total training time.

```python
eval_env_cfg = config_dict.ConfigDict(env_cfg)
eval_env_cfg.vision_config.nworld = num_eval_envs
eval_env = registry.load(env_name, config=eval_env_cfg)

make_inference_fn, params, metrics = train_fn(
    environment=env,
    wrap_env_fn=wrapper.wrap_for_brax_training,
    eval_env=eval_env,
)
print(f"time to jit: {times[1] - times[0]}")
print(f"time to train: {times[-1] - times[1]}")
```

--------------------------------

### List and Load Manipulation Environments

Source: https://context7.com/google-deepmind/mujoco_playground/llms.txt

List all available manipulation environments and load specific tasks like Panda arm pick cube or LEAP hand cube reorientation. Note the action and observation sizes for different manipulation setups.

```python
import jax
from mujoco_playground import manipulation

# List all manipulation environments
print(manipulation.ALL_ENVS)
# ('AlohaHandOver', 'AlohaSinglePegInsertion', 'PandaPickCube',
#  'PandaPickCubeOrientation', 'PandaPickCubeCartesian', 'PandaOpenCabinet',
#  'LeapCubeReorient', 'LeapCubeRotateZAxis', 'AeroCubeRotateZAxis', ...)

# Load Panda arm pick cube task
env = manipulation.load("PandaPickCube")

rng = jax.random.PRNGKey(0)
state = env.reset(rng)

print(f"Action size: {env.action_size}")  # 7 (joint velocities)
print(f"Observation size: {env.observation_size}")

# Load LEAP hand cube reorientation
env = manipulation.load("LeapCubeReorient")
state = env.reset(rng)
print(f"LEAP hand action size: {env.action_size}")  # 16 DOF

# Get domain randomizer for sim-to-real
randomizer = manipulation.get_domain_randomizer("LeapCubeReorient")
print(f"Domain randomizer available: {randomizer is not None}")  # True
```

--------------------------------

### Get TensorFlow Policy Network Output Shape

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/brax_network_to_onnx.ipynb

Provides an example of how to get the output shape of a TensorFlow policy network by passing a zero tensor as input. This is useful for verifying the network's output dimensions.

```python
example_input = tf.zeros((1, obs_size["state"][0]))
example_output = tf_policy_network(example_input)
print(example_output.shape)
```

--------------------------------

### Check MuJoCo Installation and GPU Setup

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/manipulation.ipynb

Verifies the MuJoCo installation and GPU communication. It configures MuJoCo to use the EGL rendering backend and sets XLA flags for performance optimization on GPUs. A GPU runtime is mandatory.

```python
# @title Check if MuJoCo installation was successful

import distutils.util
import os
import subprocess

if subprocess.run('nvidia-smi').returncode:
  raise RuntimeError(
      'Cannot communicate with GPU. '
      'Make sure you are using a GPU Colab runtime. '
      'Go to the Runtime menu and select Choose runtime type.'
  )

# Add an ICD config so that glvnd can pick up the Nvidia EGL driver.
# This is usually installed as part of an Nvidia driver package, but the Colab
# kernel doesn't install its driver via APT, and as a result the ICD is missing.
# (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)
NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'
if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):
  with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:
    f.write("""
{
    "file_format_version" : "1.0.0",
    "ICD" : {
        "library_path" : "libEGL_nvidia.so.0"
    }
}
"""      )

# Configure MuJoCo to use the EGL rendering backend (requires GPU)
print('Setting environment variable to use GPU rendering:')
%env MUJOCO_GL=egl

try:
  print('Checking that the installation succeeded:')
  import mujoco

  mujoco.MjModel.from_xml_string('<mujoco/>')
except Exception as e:
  raise e from RuntimeError(
      'Something went wrong during installation. Check the shell output above '
      'for more information.\n'
      'If using a hosted Colab runtime, make sure you enable GPU acceleration '
      'by going to the Runtime menu and selecting "Choose runtime type".'
  )

print('Installation successful.')

# Tell XLA to use Triton GEMM, this improves steps/sec by ~30% on some GPUs
xla_flags = os.environ.get('XLA_FLAGS', '')
xla_flags += ' --xla_gpu_triton_gemm_any=True'
os.environ['XLA_FLAGS'] = xla_flags
```

--------------------------------

### Configure Environment and Training Parameters

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/getup.ipynb

Sets up the environment name, retrieves default environment configuration, initializes domain randomization, and loads PPO training parameters.

```python
env_name = "Go1Getup"
env_cfg = locomotion.get_default_config(env_name)
randomizer = locomotion.get_domain_randomizer(env_name)
ppo_params = locomotion_params.brax_ppo_config(env_name)
```

--------------------------------

### Verify MuJoCo Playground Installation

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/README.md

Verify the installation of MuJoCo Playground by running a simple Python command.

```sh
uv --no-config run python -c "import mujoco_playground; print('Success')"
```

--------------------------------

### Configure Environment and Training Parameters

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/handstand.ipynb

Sets up the environment name, retrieves default environment configuration, obtains a domain randomizer, and loads PPO training parameters.

```python
env_name = "Go1Footstand"
env_cfg = locomotion.get_default_config(env_name)
randomizer = locomotion.get_domain_randomizer(env_name)
ppo_params = locomotion_params.brax_ppo_config(env_name)
```

--------------------------------

### Install HIDAPI and Set DYLD_LIBRARY_PATH on macOS

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/sim2sim/README.md

Install the HIDAPI library using Homebrew and configure the DYLD_LIBRARY_PATH environment variable for macOS.

```bash
brew install hidapi
export DYLD_LIBRARY_PATH=/opt/homebrew/Cellar/hidapi/0.14.0/lib:$DYLD_LIBRARY_PATH
```

--------------------------------

### Configure and Execute PPO Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/handstand.ipynb

Sets up the PPO training function with specified parameters, including network factory, checkpoint restoration, progress callback, environment wrapper, randomization, and policy parameter saving. It then loads the environment, runs the training, and prints timing information.

```python
training_params = dict(ppo_params)
del training_params["network_factory"]

train_fn = functools.partial(
  ppo.train,
  **training_params,
  network_factory=functools.partial(
      ppo_networks.make_ppo_networks,
      **ppo_params.network_factory
  ),
  restore_checkpoint_path=restore_checkpoint_path,
  progress_fn=progress,
  wrap_env_fn=wrapper.wrap_for_brax_training,
  policy_params_fn=policy_params_fn,
  randomization_fn=randomizer,
)

env = locomotion.load(env_name, config=env_cfg)
eval_env = locomotion.load(env_name, config=env_cfg)
make_inference_fn, params, _ = train_fn(environment=env, eval_env=eval_env)
if len(times) > 1:
  print(f"time to jit: {times[1] - times[0]}")
  print(f"time to train: {times[-1] - times[1]}")

# Make a final plot of reward and success vs WALLCLOCK time.
plt.figure()
plt.ylim([0, 40])
plt.xlabel("wallclock time (s)")
plt.ylabel("reward per episode")
plt.title(f"y={y_data[-1]:.3f}")
plt.errorbar(
    [(t - times[0]).total_seconds() for t in times[:-1] ],
    y_data,
    yerr=y_dataerr,
    color="blue",
)
plt.show()
```

--------------------------------

### Define Environment and Training Parameters

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/cube_reorient.ipynb

Sets up the environment name, retrieves default configuration and domain randomizer, and defines PPO training parameters. This is a foundational step for running experiments.

```python
env_name = "LeapCubeReorient"
env_cfg = manipulation.get_default_config(env_name)
randomizer = manipulation.get_domain_randomizer(env_name)
ppo_params = manipulation_params.brax_ppo_config(env_name)
```

--------------------------------

### Install MuJoCo Playground

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/manipulation.ipynb

Installs the MuJoCo Playground package using pip. This package provides wrappers and registry functionalities for MuJoCo environments.

```python
#@title Install MuJoCo Playground
!pip install playground
```

--------------------------------

### Visualize Training Trajectories with rscope

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/README.md

Launch a training run and visualize trajectories interactively using rscope. Requires rscope to be installed (`pip install rscope`).

```python
python learning/train_jax_ppo.py --env_name PandaPickCube --rscope_envs 16 --run_evals=False --deterministic_rscope=True
```

--------------------------------

### List All Manipulation Environments

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/manipulation.ipynb

Use this to get a list of all available manipulation environments in the registry.

```python
registry.manipulation.ALL_ENVS
```

--------------------------------

### Configure and Execute PPO Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/cube_rotate.ipynb

Sets up the PPO training function with specified parameters, network factory, checkpoint restoration, progress callback, environment wrappers, randomization, and then executes the training.

```python
training_params = dict(ppo_params)
del training_params["network_factory"]

train_fn = functools.partial(
  ppo.train,
  **training_params,
  network_factory=functools.partial(
      ppo_networks.make_ppo_networks,
      **ppo_params.network_factory
  ),
  restore_checkpoint_path=restore_checkpoint_path,
  progress_fn=progress,
  wrap_env_fn=wrapper.wrap_for_brax_training,
  policy_params_fn=policy_params_fn,
  randomization_fn=randomizer,
)

env = manipulation.load(env_name, config=env_cfg)
eval_env = manipulation.load(env_name, config=env_cfg)
make_inference_fn, params, _ = train_fn(environment=env, eval_env=eval_env)
if len(times) > 1:
  print(f"time to jit: {times[1] - times[0]}")
  print(f"time to train: {times[-1] - times[1]}")

# Make a final plot of reward and success vs WALLCLOCK time.
plt.figure()
# plt.ylim([0, 75])
plt.xlabel("wallclock time (s)")
plt.ylabel("reward per episode")
plt.title(f"y={y_data[-1]:.3f}")
plt.errorbar(
    [(t - times[0]).total_seconds() for t in times[:-1]],
    y_data,
    yerr=y_dataerr,
    color="blue",
)
plt.show()
```

--------------------------------

### Prepare and Execute PPO Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb

Sets up and initiates the PPO training process, incorporating domain randomization and a progress callback. This function returns the inference function, trained parameters, and final metrics.

```python
import functools
from brax import ppo
from brax.training import ppo_networks
from mujoco_playground.utils import wrapper

randomizer = registry.get_domain_randomizer(env_name)
ppo_training_params = dict(ppo_params)
network_factory = ppo_networks.make_ppo_networks
if "network_factory" in ppo_params:
  del ppo_training_params["network_factory"]
  network_factory = functools.partial(
      ppo_networks.make_ppo_networks,
      **ppo_params.network_factory
  )

train_fn = functools.partial(
    ppo.train, **dict(ppo_training_params),
    network_factory=network_factory,
    randomization_fn=randomizer,
    progress_fn=progress
)

make_inference_fn, params, metrics = train_fn(
    environment=env,
    eval_env=registry.load(env_name, config=env_cfg),
    wrap_env_fn=wrapper.wrap_for_brax_training,
)
print(f"time to jit: {times[1] - times[0]}")
print(f"time to train: {times[-1] - times[1]}")
```

--------------------------------

### Prepare Inference Environment and Functions

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/vision.ipynb

Sets up the inference environment with a specified number of worlds and wraps it for Brax training. It then JIT-compiles the reset, step, and inference functions for efficient execution.

```python
infer_env_cfg = config_dict.ConfigDict(env_cfg)
infer_env_cfg.vision_config.nworld = 64
infer_env = registry.load(env_name, config=infer_env_cfg)
wrapped_infer_env = wrapper.wrap_for_brax_training(
    infer_env,
    episode_length=ppo_params.episode_length,
    action_repeat=ppo_params.get("action_repeat", 1),
)

jit_reset = jax.jit(wrapped_infer_env.reset)
jit_step = jax.jit(wrapped_infer_env.step)
jit_inference_fn = jax.jit(make_inference_fn(params, deterministic=True))
```

--------------------------------

### Get Default Environment Configuration

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/dm_control_suite.ipynb

Retrieves the default configuration for a specified MuJoCo Playground environment. This config can be modified.

```python
env_cfg = registry.get_default_config('CartpoleBalance')
env_cfg
```

--------------------------------

### Get JAX Device Information

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/ablate_device_topology.ipynb

Retrieves the number and type of available JAX devices, formatting them into a descriptive string.

```python
num_devices = len(jax.devices())
device_kind = jax.devices()[0].device_kind
device_topo = f'{num_devices}x {device_kind}'
```

--------------------------------

### Get Domain Randomizer for Environment

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb

Retrieve the domain randomizer function for the specified environment. This can be used to improve policy robustness.

```python
registry.get_domain_randomizer(env_name)
```

--------------------------------

### Get PPO and SAC Training Parameters

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/dm_control_suite.ipynb

Retrieves default configuration parameters for PPO and SAC training from the dm_control_suite_params module.

```python
from mujoco_playground.config import dm_control_suite_params

ppo_params = dm_control_suite_params.brax_ppo_config(env_name)
sac_params = dm_control_suite_params.brax_sac_config(env_name)
```

--------------------------------

### Initialize Environment and JIT Functions

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/viz_sim_policies.ipynb

Loads a specific environment ('CheetahRun') and its default configuration, then JIT-compiles the reset and step functions for performance.

```python
env_name = 'CheetahRun'
N_EPISODES = 1

env = registry.load(env_name)
cfg = registry.get_default_config(env_name)

jit_reset = jax.jit(env.reset)
jit_step = jax.jit(env.step)
```

--------------------------------

### Create Inference Function

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/brax_network_to_onnx.ipynb

Generates a callable inference function from the PPO network and its parameters. This function can be used to get actions from observations.

```python
make_inference_fn = ppo_networks.make_inference_fn(ppo_network)
inference_fn = make_inference_fn(params, deterministic=True)
```

--------------------------------

### Initialize Environment Configuration

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/position_cube.ipynb

Sets the environment name and retrieves its default configuration. This is the first step in setting up a new experiment.

```python
env_name = "PandaOrientation"
env_cfg = manipulation.get_default_config(env_name)
```

--------------------------------

### Import Necessary Libraries for Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/apollo_joystick.ipynb

Imports essential libraries for JAX, Brax, MuJoCo, Orbax checkpointing, and plotting. Ensure these libraries are installed before running.

```python
import functools
import json
import pickle
from datetime import datetime

import jax
import mediapy as media
import mujoco
import numpy as np
from brax.training.agents.ppo import networks as ppo_networks
from brax.training.agents.ppo import train as ppo
from etils import epath
from flax.training import orbax_utils
from orbax import checkpoint as ocp

from mujoco_playground import registry, wrapper
from mujoco_playground.config import locomotion_params
from mujoco_playground.experimental.utils.plotting import TrainingPlotter

# Enable persistent compilation cache.
jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
```

--------------------------------

### Load Unitree Go1 Joystick Environment

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb

Load the Unitree Go1 joystick environment and its default configuration. This is a starting point for training locomotion policies.

```python
env_name = 'Go1JoystickFlatTerrain'
env = registry.load(env_name)
env_cfg = registry.get_default_config(env_name)
```

--------------------------------

### Run Training with JAX PPO

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/README.md

Run the JAX PPO training script with a specified environment. This is the basic usage command.

```bash
train-jax-ppo --env_name CartpoleBalance
```

--------------------------------

### List and Load Environments with Registry

Source: https://context7.com/google-deepmind/mujoco_playground/llms.txt

Use the registry module to list all available environments and load them with default or custom configurations. Configuration overrides can specify parameters like control time step, episode length, and the backend implementation.

```python
import jax
from mujoco_playground import registry

# List all available environments
print(registry.ALL_ENVS)
# ('AcrobotSwingup', 'CartpoleBalance', 'CheetahRun', 'Go1JoystickFlatTerrain',
#  'G1JoystickFlatTerrain', 'PandaPickCube', 'LeapCubeReorient', ...)

# Load an environment with default configuration
env = registry.load("CartpoleBalance")

# Load with configuration overrides
env = registry.load(
    "Go1JoystickFlatTerrain",
    config_overrides={
        "ctrl_dt": 0.02,
        "episode_length": 500,
        "impl": "warp"  # Use MuJoCo Warp backend
    }
)

# Get default configuration for an environment
config = registry.get_default_config("CartpoleBalance")
print(config)
# ctrl_dt: 0.01
# sim_dt: 0.01
# episode_length: 1000
# action_repeat: 1
# impl: 'warp'
```

--------------------------------

### Get Brax PPO Configuration for Go1

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb

Retrieve the default PPO training parameters for the specified environment. These parameters are used for training the locomotion policy.

```python
from mujoco_playground.config import locomotion_params
ppo_params = locomotion_params.brax_ppo_config(env_name)
ppo_params
```

--------------------------------

### Initialize Vision-Enabled Environment

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/vision.ipynb

Initializes a Mujoco environment with vision capabilities enabled. This setup is used for training manipulation policies that rely on visual input.

```python
env_name = "PandaPickCubeCartesian"
env_cfg = registry.get_default_config(env_name)
env_cfg["vision"] = True
env = registry.load(env_name, config=env_cfg)
```

--------------------------------

### Initialize Evaluation Environment and Functions

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/humanoid_joystick.ipynb

Sets up the evaluation environment and JIT-compiles reset and step functions for efficient interaction. Imports necessary modules for gait analysis.

```python
import jax
import jax.numpy as jp
from mujoco_playground._src.gait import draw_joystick_command
from locomotion import locomotion

# Enable perturbation in the eval env.
env_cfg = locomotion.get_default_config(env_name)
eval_env = locomotion.load(env_name, config=env_cfg)

jit_reset = jax.jit(eval_env.reset)
jit_step = jax.jit(eval_env.step)
```

--------------------------------

### Configure and Initialize PPO Training

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb

Sets up the PPO training function with environment randomization, progress reporting, and policy parameter saving. It handles network factory configurations and merges training parameters.

```python
randomizer = registry.get_domain_randomizer(env_name)
ppo_training_params = dict(ppo_params)
network_factory = ppo_networks.make_ppo_networks
if "network_factory" in ppo_params:
  del ppo_training_params["network_factory"]
  network_factory = functools.partial(
      ppo_networks.make_ppo_networks,
      **ppo_params.network_factory
  )

train_fn = functools.partial(
    ppo.train, **dict(ppo_training_params),
    network_factory=network_factory,
    randomization_fn=randomizer,
    progress_fn=progress,
    policy_params_fn=policy_params_fn,
)
```

--------------------------------

### Render Environment Trajectories

Source: https://context7.com/google-deepmind/mujoco_playground/llms.txt

Load an environment and prepare for rendering trajectories to video for visualization. This setup is the first step in visualizing agent behavior over time.

```python
import jax
import mediapy as media
from mujoco_playground import registry

env = registry.load("CheetahRun")

```

--------------------------------

### Initialize Inference and Environment Functions

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/cube_rotate.ipynb

Initializes the inference function and environment reset/step functions, JIT-compiling them for performance. This setup is crucial for efficient simulation execution.

```python
inference_fn = make_inference_fn(params, deterministic=True)
jit_inference_fn = jax.jit(inference_fn)

eval_env = manipulation.load(env_name, config=env_cfg)
jit_reset = jax.jit(eval_env.reset)
jit_step = jax.jit(eval_env.step)
```

--------------------------------

### Load Environment Configuration and Parameters

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/cube_rotate.ipynb

Loads the default configuration for the specified environment and retrieves domain randomization parameters and PPO training configuration.

```python
env_name = "LeapCubeRotateZAxis"
env_cfg = manipulation.get_default_config(env_name)
randomizer = manipulation.get_domain_randomizer(env_name)
ppo_params = manipulation_params.brax_ppo_config(env_name)
```

--------------------------------

### Configure XLA and Environment Variables

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/ablate_device_topology.ipynb

Sets environment variables for XLA GPU, preallocation, and MuJoCo's OpenGL context. This is typically done once at the start of a script.

```python
import os

xla_flags = os.environ.get("XLA_FLAGS", "")
xla_flags += " --xla_gpu_triton_gemm_any=True"
os.environ["XLA_FLAGS"] = xla_flags
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
os.environ["MUJOCO_GL"] = "egl"
```

--------------------------------

### Load Environment and Training Parameters

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb

Loads the environment, its default configuration, and PPO training parameters for a specified environment name. Ensure the 'registry' and 'locomotion_params' are correctly imported.

```python
from mujoco_playground.config import locomotion_params

env_name = 'Go1Handstand'
env = registry.load(env_name)
env_cfg = registry.get_default_config(env_name)
ppo_params = locomotion_params.brax_ppo_config(env_name)
```

--------------------------------

### Configure Checkpointing and Training Function

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/joystick_gait_tracking.ipynb

Sets up the checkpoint directory and defines the PPO training function with specified hyperparameters and callbacks.

```python
ckpt_path = epath.Path("checkpoints").resolve() / exp_name
ckpt_path.mkdir(parents=True, exist_ok=True)
print(f"Checkpoint path: {ckpt_path}")

with open(ckpt_path / "config.json", "w") as fp:
  json.dump(env_cfg.to_dict(), fp, indent=4)


def policy_params_fn(current_step, make_policy, params):
  orbax_checkpointer = ocp.PyTreeCheckpointer()
  save_args = orbax_utils.save_args_from_target(params)
  path = ckpt_path / f"{current_step}"
  orbax_checkpointer.save(path, params, force=True, save_args=save_args)


make_networks_factory = functools.partial(
    ppo_networks.make_ppo_networks,
    policy_hidden_layer_sizes=(128, 128, 128, 128),
)

train_fn = functools.partial(
    ppo.train,
    num_timesteps=70_000_000,
    num_evals=5,
    reward_scaling=1.0,
    episode_length=env_cfg.episode_length,
    normalize_observations=True,
    action_repeat=1,
    unroll_length=20,
    num_minibatches=32,
    num_updates_per_batch=4,
    discounting=0.97,
    learning_rate=3e-4,
    entropy_cost=1e-2,
    num_envs=8192,
    batch_size=256,
    network_factory=make_networks_factory,
    policy_params_fn=policy_params_fn,
    randomization_fn=randomizer,
    seed=1,
    restore_checkpoint_path=restore_checkpoint_path,
    max_grad_norm=1.0,
)
```

--------------------------------

### Import Plotting and Graphics Packages

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/manipulation.ipynb

Imports essential libraries for plotting and creating graphics, including mediapy and matplotlib. It also ensures ffmpeg is installed for media processing.

```python
# @title Import packages for plotting and creating graphics
import json
import itertools
import time
from typing import Callable, List, NamedTuple, Optional, Union
import numpy as np

# Graphics and plotting.
print("Installing mediapy:")
!command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)
!pip install -q mediapy
import mediapy as media
import matplotlib.pyplot as plt

# More legible printing from numpy.
np.set_printoptions(precision=3, suppress=True, linewidth=100)
```

--------------------------------

### Prepare Inference Environment and JIT Functions

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/vision.ipynb

This snippet configures the environment for inference, wraps it for Brax training, and then JIT-compiles the reset, step, and inference functions for optimized performance.

```python
infer_env_cfg = config_dict.ConfigDict(env_cfg)
infer_env_cfg.vision_config.nworld = 1
infer_env = registry.load(env_name, config=infer_env_cfg)
wrapped_infer_env = wrapper.wrap_for_brax_training(
    infer_env,
    episode_length=ppo_params.episode_length,
    action_repeat=ppo_params.get("action_repeat", 1),
)

jit_reset = jax.jit(wrapped_infer_env.reset)
jit_step = jax.jit(wrapped_infer_env.step)
jit_inference_fn = jax.jit(make_inference_fn(params, deterministic=True))
```

--------------------------------

### Import Core Libraries for MuJoCo Playground

Source: https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/experimental/learning/dm_control_suite.ipynb

Imports essential libraries for JAX, plotting, media handling, and specific MuJoCo Playground modules. Ensure these are installed before running.

```python
import functools
import re
from datetime import datetime
from typing import Tuple

import jax
import jax.numpy as jp
import matplotlib.pyplot as plt
import mediapy as media
import numpy as np
from brax.training.agents.sac import train as sac
from IPython.display import clear_output, display

from mujoco_playground import dm_control_suite as suite
from mujoco_playground import wrapper
```