{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "MpkYHwCqk7W-" }, "source": [ "![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n", "\n", "# \u003ch1\u003e\u003ccenter\u003eTutorial \u003ca href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/\u003e\u003c/a\u003e\u003c/center\u003e\u003c/h1\u003e\n", "\n", "This notebook provides an introductory tutorial for [**MuJoCo XLA (MJX)**](https://github.com/google-deepmind/mujoco/blob/main/mjx), a JAX-based implementation of MuJoCo useful for RL training workloads.\n", "\n", "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime \u003e Change runtime type\".\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n" ] }, { "cell_type": "markdown", "metadata": { "id": "xBSdkbmGN2K-" }, "source": [ "### Copyright notice" ] }, { "cell_type": "markdown", "metadata": { "id": "_UbO9uhtBSX5" }, "source": [ "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eCopyright 2023 DeepMind Technologies Limited.\u003c/small\u003e\u003c/p\u003e\n", "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eLicensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at \u003ca href=\"http://www.apache.org/licenses/LICENSE-2.0\"\u003ehttp://www.apache.org/licenses/LICENSE-2.0\u003c/a\u003e.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e\n", "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eUnless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e" ] }, { "cell_type": "markdown", "metadata": { "id": "YvyGCsgSCxHQ" }, "source": [ "# Install MuJoCo, MJX, and Brax" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Xqo7pyX-n72M" }, "outputs": [], "source": [ "!pip install mujoco\n", "!pip install mujoco_mjx\n", "!pip install brax" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "IbZxYDxzoz5R" }, "outputs": [], "source": [ "#@title Check if MuJoCo installation was successful\n", "\n", "from google.colab import files\n", "\n", "import distutils.util\n", "import os\n", "import subprocess\n", "if subprocess.run('nvidia-smi').returncode:\n", " raise RuntimeError(\n", " 'Cannot communicate with GPU. '\n", " 'Make sure you are using a GPU Colab runtime. '\n", " 'Go to the Runtime menu and select Choose runtime type.')\n", "\n", "# Add an ICD config so that glvnd can pick up the Nvidia EGL driver.\n", "# This is usually installed as part of an Nvidia driver package, but the Colab\n", "# kernel doesn't install its driver via APT, and as a result the ICD is missing.\n", "# (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)\n", "NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'\n", "if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):\n", " with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:\n", " f.write(\"\"\"{\n", " \"file_format_version\" : \"1.0.0\",\n", " \"ICD\" : {\n", " \"library_path\" : \"libEGL_nvidia.so.0\"\n", " }\n", "}\n", "\"\"\")\n", "\n", "# Configure MuJoCo to use the EGL rendering backend (requires GPU)\n", "print('Setting environment variable to use GPU rendering:')\n", "%env MUJOCO_GL=egl\n", "\n", "try:\n", " print('Checking that the installation succeeded:')\n", " import mujoco\n", " mujoco.MjModel.from_xml_string('\u003cmujoco/\u003e')\n", "except Exception as e:\n", " raise e from RuntimeError(\n", " 'Something went wrong during installation. Check the shell output above '\n", " 'for more information.\\n'\n", " 'If using a hosted Colab runtime, make sure you enable GPU acceleration '\n", " 'by going to the Runtime menu and selecting \"Choose runtime type\".')\n", "\n", "print('Installation successful.')" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "T5f4w3Kq2X14" }, "outputs": [], "source": [ "#@title Import packages for plotting and creating graphics\n", "import time\n", "import itertools\n", "import numpy as np\n", "from typing import Callable, NamedTuple, Optional, Union, List\n", "\n", "# Graphics and plotting.\n", "print('Installing mediapy:')\n", "!command -v ffmpeg \u003e/dev/null || (apt update \u0026\u0026 apt install -y ffmpeg)\n", "!pip install -q mediapy\n", "import mediapy as media\n", "import matplotlib.pyplot as plt\n", "\n", "# More legible printing from numpy.\n", "np.set_printoptions(precision=3, suppress=True, linewidth=100)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "ObF1UXrkb0Nd" }, "outputs": [], "source": [ "#@title Import MuJoCo, MJX, and Brax\n", "\n", "from datetime import datetime\n", "import functools\n", "import jax\n", "from jax import numpy as jp\n", "import numpy as np\n", "from typing import Any, Dict, Tuple, Union\n", "\n", "from brax import envs\n", "from brax import math\n", "from brax.base import Base, Motion, Transform\n", "from brax.envs.base import Env, State\n", "from brax.training.agents.ppo import train as ppo\n", "from brax.training.agents.ppo import networks as ppo_networks\n", "from brax.io import model\n", "from etils import epath\n", "from flax import struct\n", "from matplotlib import pyplot as plt\n", "import mediapy as media\n", "from ml_collections import config_dict\n", "import mujoco\n", "from mujoco import mjx\n" ] }, { "cell_type": "markdown", "metadata": { "id": "RAv6WUVUm78k" }, "source": [ "# Training a Policy with MJX\n", "MJX is an implementation of MuJoCo written in [JAX](https://jax.readthedocs.io/en/latest/index.html), enabling large batch training on GPU/TPU. In this notebook, we demonstrate how to train RL policies with MJX.\n", "\n", "First, we implement an environment `State` so that we can plug into the [Brax](https://github.com/google/brax) environment API. `State` holds the observation, reward, metrics, and environment info. Notably `State.pipeline_state` holds a `mjx.Data` object, which is analogous to `mjData` in MuJoCo.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "7DQ_rW4CkIB_" }, "outputs": [], "source": [ "#@title State\n", "\n", "@struct.dataclass\n", "class State(Base):\n", " \"\"\"Environment state for training and inference with brax.\n", "\n", " Args:\n", " pipeline_state: the physics state, mjx.Data\n", " obs: environment observations\n", " reward: environment reward\n", " done: boolean, True if the current episode has terminated\n", " metrics: metrics that get tracked per environment step\n", " info: environment variables defined and updated by the environment reset\n", " and step functions\n", " \"\"\"\n", "\n", " pipeline_state: mjx.Data\n", " obs: jax.Array\n", " reward: jax.Array\n", " done: jax.Array\n", " metrics: Dict[str, jax.Array] = struct.field(default_factory=dict)\n", " info: Dict[str, Any] = struct.field(default_factory=dict)\n" ] }, { "cell_type": "markdown", "metadata": { "id": "acpXtDLNXLV9" }, "source": [ "\n", "Next, we implement `MjxEnv`, an environment class we'll use through the notebook. `MjxEnv` initializes a `mjx.Model` and `mjx.Data` object. Notice that `MjxEnv` calls `mjx.step` for every `pipeline_step`, which is analgous to `mujoco.mj_step`.\n", "\n", "`MjxEnv` also inherits from `brax.envs.base.Env` which allows us to use the training agents implemented in brax." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "ccujYeJ5XOhx" }, "outputs": [], "source": [ "#@title MjxEnv\n", "\n", "class MjxEnv(Env):\n", " \"\"\"API for driving an MJX system for training and inference in brax.\"\"\"\n", "\n", " def __init__(\n", " self,\n", " mj_model: mujoco.MjModel,\n", " physics_steps_per_control_step: int = 1,\n", " ):\n", " \"\"\"Initializes MjxEnv.\n", "\n", " Args:\n", " mj_model: mujoco.MjModel\n", " physics_steps_per_control_step: the number of times to step the physics\n", " pipeline for each environment step\n", " \"\"\"\n", " self.model = mj_model\n", " self.data = mujoco.MjData(mj_model)\n", " self.sys = mjx.device_put(mj_model)\n", " self._physics_steps_per_control_step = physics_steps_per_control_step\n", "\n", " def pipeline_init(\n", " self, qpos: jax.Array, qvel: jax.Array\n", " ) -\u003e mjx.Data:\n", " \"\"\"Initializes the physics state.\"\"\"\n", " data = mjx.device_put(self.data)\n", " data = data.replace(qpos=qpos, qvel=qvel, ctrl=jp.zeros(self.sys.nu))\n", " data = mjx.forward(self.sys, data)\n", " return data\n", "\n", " def pipeline_step(\n", " self, data: mjx.Data, ctrl: jax.Array\n", " ) -\u003e mjx.Data:\n", " \"\"\"Takes a physics step using the physics pipeline.\"\"\"\n", " def f(data, _):\n", " data = data.replace(ctrl=ctrl)\n", " return (\n", " mjx.step(self.sys, data),\n", " None,\n", " )\n", " data, _ = jax.lax.scan(f, data, (), self._physics_steps_per_control_step)\n", " return data\n", "\n", " @property\n", " def dt(self) -\u003e jax.Array:\n", " \"\"\"The timestep used for each env step.\"\"\"\n", " return self.sys.opt.timestep * self._physics_steps_per_control_step\n", "\n", " @property\n", " def observation_size(self) -\u003e int:\n", " rng = jax.random.PRNGKey(0)\n", " reset_state = self.unwrapped.reset(rng)\n", " return reset_state.obs.shape[-1]\n", "\n", " @property\n", " def action_size(self) -\u003e int:\n", " return self.sys.nu\n", "\n", " @property\n", " def backend(self) -\u003e str:\n", " return 'mjx'\n", "\n", " def _pos_vel(\n", " self, data: mjx.Data\n", " ) -\u003e Tuple[Transform, Motion]:\n", " \"\"\"Returns 6d spatial transform and 6d velocity for all bodies.\"\"\"\n", " x = Transform(pos=data.xpos[1:, :], rot=data.xquat[1:, :])\n", " cvel = Motion(vel=data.cvel[1:, 3:], ang=data.cvel[1:, :3])\n", " offset = data.xpos[1:, :] - data.subtree_com[\n", " self.model.body_rootid[np.arange(1, self.model.nbody)]]\n", " xd = Transform.create(pos=offset).vmap().do(cvel)\n", " return x, xd\n" ] }, { "cell_type": "markdown", "metadata": { "id": "iPlFu4CiIgBN" }, "source": [ "Finally we can implement a real environment. We choose to first implement the Humanoid environment. Notice that `reset` initializes a `State`, and `step` steps through the physics step and reward logic. The reward and stepping logic train the Humanoid to run forwards." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "mtGMYNLE3QJN" }, "outputs": [], "source": [ "#@title Humanoid Env\n", "\n", "class Humanoid(MjxEnv):\n", "\n", " def __init__(\n", " self,\n", " forward_reward_weight=1.25,\n", " ctrl_cost_weight=0.1,\n", " healthy_reward=5.0,\n", " terminate_when_unhealthy=True,\n", " healthy_z_range=(1.0, 2.0),\n", " reset_noise_scale=1e-2,\n", " exclude_current_positions_from_observation=True,\n", " **kwargs,\n", " ):\n", " path = epath.Path(epath.resource_path('mujoco')) / (\n", " 'mjx/benchmark/model/humanoid'\n", " )\n", " mj_model = mujoco.MjModel.from_xml_path(\n", " (path / 'humanoid.xml').as_posix())\n", " mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n", " mj_model.opt.iterations = 6\n", " mj_model.opt.ls_iterations = 6\n", "\n", " physics_steps_per_control_step = 5\n", " kwargs['physics_steps_per_control_step'] = kwargs.get(\n", " 'physics_steps_per_control_step', physics_steps_per_control_step)\n", "\n", " super().__init__(mj_model=mj_model, **kwargs)\n", "\n", " self._forward_reward_weight = forward_reward_weight\n", " self._ctrl_cost_weight = ctrl_cost_weight\n", " self._healthy_reward = healthy_reward\n", " self._terminate_when_unhealthy = terminate_when_unhealthy\n", " self._healthy_z_range = healthy_z_range\n", " self._reset_noise_scale = reset_noise_scale\n", " self._exclude_current_positions_from_observation = (\n", " exclude_current_positions_from_observation\n", " )\n", "\n", " def reset(self, rng: jp.ndarray) -\u003e State:\n", " \"\"\"Resets the environment to an initial state.\"\"\"\n", " rng, rng1, rng2 = jax.random.split(rng, 3)\n", "\n", " low, hi = -self._reset_noise_scale, self._reset_noise_scale\n", " qpos = self.sys.qpos0 + jax.random.uniform(\n", " rng1, (self.sys.nq,), minval=low, maxval=hi\n", " )\n", " qvel = jax.random.uniform(\n", " rng2, (self.sys.nv,), minval=low, maxval=hi\n", " )\n", "\n", " data = self.pipeline_init(qpos, qvel)\n", "\n", " obs = self._get_obs(data, jp.zeros(self.sys.nu))\n", " reward, done, zero = jp.zeros(3)\n", " metrics = {\n", " 'forward_reward': zero,\n", " 'reward_linvel': zero,\n", " 'reward_quadctrl': zero,\n", " 'reward_alive': zero,\n", " 'x_position': zero,\n", " 'y_position': zero,\n", " 'distance_from_origin': zero,\n", " 'x_velocity': zero,\n", " 'y_velocity': zero,\n", " }\n", " return State(data, obs, reward, done, metrics)\n", "\n", " def step(self, state: State, action: jp.ndarray) -\u003e State:\n", " \"\"\"Runs one timestep of the environment's dynamics.\"\"\"\n", " data0 = state.pipeline_state\n", " data = self.pipeline_step(data0, action)\n", "\n", " com_before = data0.subtree_com[1]\n", " com_after = data.subtree_com[1]\n", " velocity = (com_after - com_before) / self.dt\n", " forward_reward = self._forward_reward_weight * velocity[0]\n", "\n", " min_z, max_z = self._healthy_z_range\n", " is_healthy = jp.where(data.qpos[2] \u003c min_z, x=0.0, y=1.0)\n", " is_healthy = jp.where(\n", " data.qpos[2] \u003e max_z, x=0.0, y=is_healthy\n", " )\n", " if self._terminate_when_unhealthy:\n", " healthy_reward = self._healthy_reward\n", " else:\n", " healthy_reward = self._healthy_reward * is_healthy\n", "\n", " ctrl_cost = self._ctrl_cost_weight * jp.sum(jp.square(action))\n", "\n", " obs = self._get_obs(data, action)\n", " reward = forward_reward + healthy_reward - ctrl_cost\n", " done = 1.0 - is_healthy if self._terminate_when_unhealthy else 0.0\n", " state.metrics.update(\n", " forward_reward=forward_reward,\n", " reward_linvel=forward_reward,\n", " reward_quadctrl=-ctrl_cost,\n", " reward_alive=healthy_reward,\n", " x_position=com_after[0],\n", " y_position=com_after[1],\n", " distance_from_origin=jp.linalg.norm(com_after),\n", " x_velocity=velocity[0],\n", " y_velocity=velocity[1],\n", " )\n", "\n", " return state.replace(\n", " pipeline_state=data, obs=obs, reward=reward, done=done\n", " )\n", "\n", " def _get_obs(\n", " self, data: mjx.Data, action: jp.ndarray\n", " ) -\u003e jp.ndarray:\n", " \"\"\"Observes humanoid body position, velocities, and angles.\"\"\"\n", " position = data.qpos\n", " if self._exclude_current_positions_from_observation:\n", " position = position[2:]\n", "\n", " # external_contact_forces are excluded\n", " return jp.concatenate([\n", " position,\n", " data.qvel,\n", " data.cinert[1:].ravel(),\n", " data.cvel[1:].ravel(),\n", " data.qfrc_actuator,\n", " ])\n", "\n", "\n", "envs.register_environment('humanoid', Humanoid)" ] }, { "cell_type": "markdown", "metadata": { "id": "P1K6IznI2y83" }, "source": [ "## Visualize a Rollout\n", "\n", "Let's instantiate the environment and visualize a short rollout.\n", "\n", "NOTE: Since episodes terminates early if the torso is below the healthy z-range, the only relevant contacts for this task are between the feet and the plane. We turn off other contacts." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "EhKLFK54C1CH" }, "outputs": [], "source": [ "# instantiate the environment\n", "env_name = 'humanoid'\n", "env = envs.get_environment(env_name)\n", "\n", "# define the jit reset/step functions\n", "jit_reset = jax.jit(env.reset)\n", "jit_step = jax.jit(env.step)\n", "\n", "# instantiate the renderer\n", "renderer = mujoco.Renderer(env.model)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "9f2ME2WbA5Ip" }, "outputs": [], "source": [ "#@title Define a render utility function\n", "\n", "def get_image(state: State, camera: str) -\u003e np.ndarray:\n", " \"\"\"Renders the environment state.\"\"\"\n", " d = mujoco.MjData(env.model)\n", " # write the mjx.Data into an mjData object\n", " mjx.device_get_into(d, state.pipeline_state)\n", " mujoco.mj_forward(env.model, d)\n", " # use the mjData object to update the renderer\n", " renderer.update_scene(d, camera=camera)\n", " return renderer.render()\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Ph8u-v2Q2xLS" }, "outputs": [], "source": [ "# initialize the state\n", "state = jit_reset(jax.random.PRNGKey(0))\n", "rollout = [state]\n", "images = [get_image(state, camera='side')]\n", "\n", "# grab a trajectory\n", "for i in range(10):\n", " ctrl = -0.1 * jp.ones(env.sys.nu)\n", " state = jit_step(state, ctrl)\n", " rollout.append(state)\n", " images.append(get_image(state, camera='side'))\n", "\n", "media.show_video(images, fps=1.0 / env.dt)" ] }, { "cell_type": "markdown", "metadata": { "id": "BQDG6NQ1CbZD" }, "source": [ "## Train Humanoid Policy\n", "\n", "Let's finally train a policy with PPO to make the Humanoid run forwards. Training takes about 13-14 minutes on a Tesla V100 GPU." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "xLiddQYPApBw" }, "outputs": [], "source": [ "train_fn = functools.partial(\n", " ppo.train, num_timesteps=30_000_000, num_evals=5, reward_scaling=0.1,\n", " episode_length=1000, normalize_observations=True, action_repeat=1,\n", " unroll_length=10, num_minibatches=32, num_updates_per_batch=8,\n", " discounting=0.97, learning_rate=3e-4, entropy_cost=1e-3, num_envs=2048,\n", " batch_size=1024, seed=0)\n", "\n", "\n", "x_data = []\n", "y_data = []\n", "ydataerr = []\n", "times = [datetime.now()]\n", "\n", "max_y, min_y = 13000, 0\n", "def progress(num_steps, metrics):\n", " times.append(datetime.now())\n", " x_data.append(num_steps)\n", " y_data.append(metrics['eval/episode_reward'])\n", " ydataerr.append(metrics['eval/episode_reward_std'])\n", "\n", " plt.xlim([0, train_fn.keywords['num_timesteps'] * 1.25])\n", " plt.ylim([min_y, max_y])\n", "\n", " plt.xlabel('# environment steps')\n", " plt.ylabel('reward per episode')\n", " plt.title(f'y={y_data[-1]:.3f}')\n", "\n", " plt.errorbar(\n", " x_data, y_data, yerr=ydataerr)\n", " plt.show()\n", "\n", "make_inference_fn, params, _= train_fn(environment=env, progress_fn=progress)\n", "\n", "print(f'time to jit: {times[1] - times[0]}')\n", "print(f'time to train: {times[-1] - times[1]}')" ] }, { "cell_type": "markdown", "metadata": { "id": "YYIch0HEApBx" }, "source": [ "## Save and Load Policy\n", "\n", "We can save and load the policy using the brax model API." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Z8gI6qH6ApBx" }, "outputs": [], "source": [ "#@title Save Model\n", "model_path = '/tmp/mjx_brax_policy'\n", "model.save_params(model_path, params)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "h4reaWgxApBx" }, "outputs": [], "source": [ "#@title Load Model and Define Inference Function\n", "params = model.load_params(model_path)\n", "\n", "inference_fn = make_inference_fn(params)\n", "jit_inference_fn = jax.jit(inference_fn)" ] }, { "cell_type": "markdown", "metadata": { "id": "0G357XIfApBy" }, "source": [ "## Visualize Policy\n", "\n", "Finally we can visualize the policy." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "osYasMw4ApBy" }, "outputs": [], "source": [ "eval_env = envs.get_environment(env_name)\n", "\n", "jit_reset = jax.jit(eval_env.reset)\n", "jit_step = jax.jit(eval_env.step)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "d-UhypudApBy" }, "outputs": [], "source": [ "# initialize the state\n", "rng = jax.random.PRNGKey(0)\n", "state = jit_reset(rng)\n", "rollout = [state]\n", "images = [get_image(state, camera='side')]\n", "\n", "# grab a trajectory\n", "n_steps = 500\n", "render_every = 2\n", "\n", "for i in range(n_steps):\n", " act_rng, rng = jax.random.split(rng)\n", " ctrl, _ = jit_inference_fn(state.obs, act_rng)\n", " state = jit_step(state, ctrl)\n", " rollout.append(state)\n", " if i % render_every == 0:\n", " images.append(get_image(state, camera='side'))\n", "\n", " if state.done:\n", " break\n", "\n", "media.show_video(images, fps=1.0 / eval_env.dt / render_every)" ] }, { "cell_type": "markdown", "metadata": { "id": "zR-heox6LARK" }, "source": [ "# MJX Policy in MuJoCo\n", "\n", "Note that we can also perform the physics step using the original MuJoCo python bindings to show that the policy trained in MJX works in MuJoCo." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "w6ixFi4dApBy" }, "outputs": [], "source": [ "mj_model = eval_env.model\n", "mj_data = mujoco.MjData(mj_model)\n", "\n", "renderer = mujoco.Renderer(mj_model)\n", "ctrl = jp.zeros(mj_model.nu)\n", "\n", "images = []\n", "for i in range(n_steps):\n", " act_rng, rng = jax.random.split(rng)\n", "\n", " obs = eval_env._get_obs(mjx.device_put(mj_data), ctrl)\n", " ctrl, _ = jit_inference_fn(obs, act_rng)\n", "\n", " mj_data.ctrl = ctrl\n", " for _ in range(eval_env._physics_steps_per_control_step):\n", " mujoco.mj_step(mj_model, mj_data) # Physics step using MuJoCo mj_step.\n", "\n", " if i % render_every == 0:\n", " renderer.update_scene(mj_data, camera='side')\n", " images.append(renderer.render())\n", "\n", "media.show_video(images, fps=1.0 / eval_env.dt / render_every)" ] }, { "cell_type": "markdown", "metadata": { "id": "65mIPj6DQNNa" }, "source": [ "# Domain Randomization\n", "\n", "We might also want to include randomization over certain `mjModel` parameters while training a policy. In MJX, we can easily create a batch of environments with randomized values populated in `mjx.Model`. Below, we show a function that randomizes friction and actuator gain/bias." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "h8mhzKjHQuoL" }, "outputs": [], "source": [ "def domain_randomize(sys, rng):\n", " \"\"\"Randomizes the mjx.Model.\"\"\"\n", " @jax.vmap\n", " def rand(rng):\n", " _, key = jax.random.split(rng, 2)\n", " # friction\n", " friction = jax.random.uniform(key, (1,), minval=0.6, maxval=1.4)\n", " friction = sys.geom_friction.at[:, 0].set(friction)\n", " # actuator\n", " _, key = jax.random.split(key, 2)\n", " gain_range = (-10, -5)\n", " param = jax.random.uniform(\n", " key, (1,), minval=gain_range[0], maxval=gain_range[1]\n", " ) + sys.actuator_gainprm[:, 0]\n", " gain = sys.actuator_gainprm.at[:, 0].set(param)\n", " bias = sys.actuator_biasprm.at[:, 1].set(-param)\n", " return friction, gain, bias\n", "\n", " friction, gain, bias = rand(rng)\n", "\n", " in_axes = jax.tree_map(lambda x: None, sys)\n", " in_axes = in_axes.tree_replace({\n", " 'geom_friction': 0,\n", " 'actuator_gainprm': 0,\n", " 'actuator_biasprm': 0,\n", " })\n", "\n", " sys = sys.tree_replace({\n", " 'geom_friction': friction,\n", " 'actuator_gainprm': gain,\n", " 'actuator_biasprm': bias,\n", " })\n", "\n", " return sys, in_axes" ] }, { "cell_type": "markdown", "metadata": { "id": "gnsZo-GWSYYj" }, "source": [ "If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`, which returns a batched `mjModel` along with a dictionary specifying the axes that are batched." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "1K45Kp2ASV9s" }, "outputs": [], "source": [ "rng = jax.random.PRNGKey(0)\n", "rng = jax.random.split(rng, 10)\n", "batched_sys, _ = domain_randomize(env.sys, rng)\n", "\n", "print('Single env friction shape: ', env.sys.geom_friction.shape)\n", "print('Batched env friction shape: ', batched_sys.geom_friction.shape)\n", "\n", "print('Friction on geom 0: ', env.sys.geom_friction[0, 0])\n", "print('Random frictions on geom 0: ', batched_sys.geom_friction[:, 0, 0])" ] }, { "cell_type": "markdown", "metadata": { "id": "efnxNOnpQFuC" }, "source": [ "## Quadruped Env\n", "\n", "Let's define a quadruped environment that takes advantage of the domain randomization function. Here we use the [Barkour v0 Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_v0) and an environment that trains a joystick policy." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "y79PoJOCIl-O" }, "outputs": [], "source": [ "#@title Barkour v0 Quadruped Env\n", "\n", "def get_config():\n", " \"\"\"Returns reward config for barkour quadruped environment.\"\"\"\n", "\n", " def get_default_rewards_config():\n", " default_config = config_dict.ConfigDict(\n", " dict(\n", " # The coefficients for all reward terms used for training. All\n", " # physical quantities are in SI units, if no otherwise specified,\n", " # i.e. joint positions are in rad, positions are measured in meters,\n", " # torques in Nm, and time in seconds, and forces in Newtons.\n", " scales=config_dict.ConfigDict(\n", " dict(\n", " # Tracking rewards are computed using exp(-delta^2/sigma)\n", " # sigma can be a hyperparameters to tune.\n", " # Track the base x-y velocity (no z-velocity tracking.)\n", " tracking_lin_vel=1.5,\n", " # Track the angular velocity along z-axis, i.e. yaw rate.\n", " tracking_ang_vel=0.8,\n", " # Below are regularization terms, we roughly divide the\n", " # terms to base state regularizations, joint\n", " # regularizations, and other behavior regularizations.\n", " # Penalize the base velocity in z direction, L2 penalty.\n", " lin_vel_z=-2.0,\n", " # Penalize the base roll and pitch rate. L2 penalty.\n", " ang_vel_xy=-0.05,\n", " # Penalize non-zero roll and pitch angles. L2 penalty.\n", " orientation=-5.0,\n", " # L2 regularization of joint torques, |tau|^2.\n", " # torques=-0.0002,\n", " torques=-0.002,\n", " # Penalize the change in the action and encourage smooth\n", " # actions. L2 regularization |action - last_action|^2\n", " action_rate=-0.1,\n", " # Encourage long swing steps. However, it does not\n", " # encourage high clearances.\n", " feet_air_time=0.2,\n", " # Encourage no motion at zero command, L2 regularization\n", " # |q - q_default|^2.\n", " stand_still=-0.5,\n", " # Early termination penalty.\n", " termination=-1.0,\n", " # Penalizing foot slipping on the ground.\n", " foot_slip=-0.1,\n", " )\n", " ),\n", " # Tracking reward = exp(-error^2/sigma).\n", " tracking_sigma=0.25,\n", " )\n", " )\n", " return default_config\n", "\n", " default_config = config_dict.ConfigDict(\n", " dict(rewards=get_default_rewards_config(),))\n", "\n", " return default_config\n", "\n", "\n", "class BarkourEnv(MjxEnv):\n", " \"\"\"Environment for training the barkour quadruped joystick policy in MJX.\"\"\"\n", "\n", " def __init__(\n", " self,\n", " obs_noise: float = 0.05,\n", " action_scale: float=0.3,\n", " **kwargs,\n", " ):\n", " path = epath.Path(epath.resource_path('mujoco')) / (\n", " 'mjx/benchmark/model/barkour_v0/assets'\n", " )\n", " mj_model = mujoco.MjModel.from_xml_path(\n", " (path / 'barkour_v0_mjx.xml').as_posix())\n", " mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n", " mj_model.opt.iterations = 4\n", " mj_model.opt.ls_iterations = 6\n", "\n", " physics_steps_per_control_step = 10\n", " kwargs['physics_steps_per_control_step'] = kwargs.get(\n", " 'physics_steps_per_control_step', physics_steps_per_control_step)\n", " super().__init__(mj_model=mj_model, **kwargs)\n", "\n", " self.torso_idx = 1\n", " self._action_scale = action_scale\n", " self._obs_noise = obs_noise\n", " self._reset_horizon = 500\n", " self._feet_index = jp.array([3, 6, 9, 12])\n", " # local positions for each foot\n", " self._feet_pos = jp.array([\n", " [-0.191284, -0.0191638, 0.013],\n", " [-0.191284, -0.0191638, -0.013],\n", " [-0.191284, -0.0191638, 0.013],\n", " [-0.191284, -0.0191638, -0.013],\n", " ])\n", " self._init_q = mj_model.keyframe('standing').qpos\n", " self._default_ap_pose = mj_model.keyframe('standing').qpos[7:]\n", " self.reward_config = get_config()\n", " self.lowers = self._default_ap_pose - jp.array([0.2, 0.8, 0.8] * 4)\n", " self.uppers = self._default_ap_pose + jp.array([0.2, 0.8, 0.8] * 4)\n", "\n", " def sample_command(self, rng: jax.Array) -\u003e jax.Array:\n", " lin_vel_x = [-0.6, 1.0] # min max [m/s]\n", " lin_vel_y = [-0.8, 0.8] # min max [m/s]\n", " ang_vel_yaw = [-0.7, 0.7] # min max [rad/s]\n", "\n", " _, key1, key2, key3 = jax.random.split(rng, 4)\n", " lin_vel_x = jax.random.uniform(\n", " key1, (1,), minval=lin_vel_x[0], maxval=lin_vel_x[1]\n", " )\n", " lin_vel_y = jax.random.uniform(\n", " key2, (1,), minval=lin_vel_y[0], maxval=lin_vel_y[1]\n", " )\n", " ang_vel_yaw = jax.random.uniform(\n", " key3, (1,), minval=ang_vel_yaw[0], maxval=ang_vel_yaw[1]\n", " )\n", " new_cmd = jp.array([lin_vel_x[0], lin_vel_y[0], ang_vel_yaw[0]])\n", " return new_cmd\n", "\n", " def reset(self, rng: jax.Array) -\u003e State:\n", " rng, key = jax.random.split(rng)\n", "\n", " qpos = jp.array(self._init_q)\n", " qvel = jp.zeros(self.model.nv)\n", " new_cmd = self.sample_command(key)\n", " data = self.pipeline_init(qpos, qvel)\n", "\n", " state_info = {\n", " 'rng': rng,\n", " 'last_act': jp.zeros(12),\n", " 'last_vel': jp.zeros(12),\n", " 'last_contact_buffer': jp.zeros((20, 4), dtype=bool),\n", " 'command': new_cmd,\n", " 'last_contact': jp.zeros(4, dtype=bool),\n", " 'feet_air_time': jp.zeros(4),\n", " 'obs_history': jp.zeros(15 * 31),\n", " 'reward_tuple': {\n", " 'tracking_lin_vel': 0.0,\n", " 'tracking_ang_vel': 0.0,\n", " 'lin_vel_z': 0.0,\n", " 'ang_vel_xy': 0.0,\n", " 'orientation': 0.0,\n", " 'torque': 0.0,\n", " 'action_rate': 0.0,\n", " 'stand_still': 0.0,\n", " 'feet_air_time': 0.0,\n", " 'foot_slip': 0.0,\n", " },\n", " 'step': 0,\n", " }\n", "\n", " x, xd = self._pos_vel(data)\n", " obs = self._get_obs(data.qpos, x, xd, state_info)\n", " reward, done = jp.zeros(2)\n", " metrics = {'total_dist': 0.0}\n", " for k in state_info['reward_tuple']:\n", " metrics[k] = state_info['reward_tuple'][k]\n", " state = State(data, obs, reward, done, metrics, state_info)\n", " return state\n", "\n", " def step(self, state: State, action: jax.Array) -\u003e State:\n", " rng, rng_noise, cmd_rng = jax.random.split(\n", " state.info['rng'], 3\n", " )\n", "\n", " # physics step\n", " cur_action = jp.array(action)\n", " action = action[:12] * self._action_scale\n", " motor_targets = jp.clip(\n", " action + self._default_ap_pose, self.lowers, self.uppers\n", " )\n", " data = self.pipeline_step(state.pipeline_state, motor_targets)\n", "\n", " # observation data\n", " x, xd = self._pos_vel(data)\n", " obs = self._get_obs(data.qpos, x, xd, state.info)\n", " obs_noise = self._obs_noise * jax.random.uniform(\n", " rng_noise, obs.shape, minval=-1, maxval=1)\n", " qpos, qvel = data.qpos, data.qvel\n", " joint_angles = qpos[7:]\n", " joint_vel = qvel[6:]\n", "\n", " # foot contact data based on z-position\n", " foot_contact = 0.017 - self._get_feet_pos_vel(x, xd)[0][:, 2]\n", " contact = foot_contact \u003e -1e-3 # a mm or less off the floor\n", " contact_filt_mm = jp.logical_or(contact, state.info['last_contact'])\n", " contact_filt_cm = jp.logical_or(\n", " foot_contact \u003e -1e-2, state.info['last_contact']\n", " )\n", " first_contact = (state.info['feet_air_time'] \u003e 0) * (contact_filt_mm)\n", " state.info['feet_air_time'] += self.dt\n", "\n", " # reward\n", " reward_tuple = {\n", " 'tracking_lin_vel': (\n", " self._reward_tracking_lin_vel(state.info['command'], x, xd)\n", " * self.reward_config.rewards.scales.tracking_lin_vel\n", " ),\n", " 'tracking_ang_vel': (\n", " self._reward_tracking_ang_vel(state.info['command'], x, xd)\n", " * self.reward_config.rewards.scales.tracking_ang_vel\n", " ),\n", " 'lin_vel_z': (\n", " self._reward_lin_vel_z(xd)\n", " * self.reward_config.rewards.scales.lin_vel_z\n", " ),\n", " 'ang_vel_xy': (\n", " self._reward_ang_vel_xy(xd)\n", " * self.reward_config.rewards.scales.ang_vel_xy\n", " ),\n", " 'orientation': (\n", " self._reward_orientation(x)\n", " * self.reward_config.rewards.scales.orientation\n", " ),\n", " 'torque': (\n", " self._reward_torques(data.qfrc_actuator)\n", " * self.reward_config.rewards.scales.torques\n", " ),\n", " 'action_rate': (\n", " self._reward_action_rate(cur_action, state.info['last_act'])\n", " * self.reward_config.rewards.scales.action_rate\n", " ),\n", " 'stand_still': (\n", " self._reward_stand_still(\n", " state.info['command'], joint_angles, self._default_ap_pose\n", " )\n", " * self.reward_config.rewards.scales.stand_still\n", " ),\n", " 'feet_air_time': (\n", " self._reward_feet_air_time(\n", " state.info['feet_air_time'],\n", " first_contact,\n", " state.info['command'],\n", " )\n", " * self.reward_config.rewards.scales.feet_air_time\n", " ),\n", " 'foot_slip': (\n", " self._reward_foot_slip(x, xd, contact_filt_cm)\n", " * self.reward_config.rewards.scales.foot_slip\n", " ),\n", " }\n", " reward = sum(reward_tuple.values())\n", " reward = jp.clip(reward * self.dt, 0.0, 10000.0)\n", "\n", " # state management\n", " state.info['last_act'] = cur_action\n", " state.info['last_vel'] = joint_vel\n", " state.info['feet_air_time'] *= ~contact_filt_mm\n", " state.info['last_contact'] = contact\n", " state.info['last_contact_buffer'] = jp.roll(\n", " state.info['last_contact_buffer'], 1, axis=0\n", " )\n", " state.info['last_contact_buffer'] = (\n", " state.info['last_contact_buffer'].at[0].set(contact)\n", " )\n", " state.info['reward_tuple'] = reward_tuple\n", " state.info['step'] += 1\n", " state.info.update(rng=rng)\n", "\n", " # resetting logic if joint limits are reached or robot is falling\n", " done = 0.0\n", " up = jp.array([0.0, 0.0, 1.0])\n", " done = jp.where(jp.dot(math.rotate(up, x.rot[0]), up) \u003c 0, 1.0, done)\n", " done = jp.where(jp.logical_or(\n", " jp.any(joint_angles \u003c .98 * self.lowers),\n", " jp.any(joint_angles \u003e .98 * self.uppers)), 1.0, done)\n", " done = jp.where(x.pos[self.torso_idx, 2] \u003c 0.18, 1.0, done)\n", "\n", " # termination reward\n", " reward += jp.where(\n", " (done == 1.0) \u0026 (state.info['step'] \u003c self._reset_horizon),\n", " self.reward_config.rewards.scales.termination,\n", " 0.0,\n", " )\n", "\n", " # when done, sample new command if more than _reset_horizon timesteps\n", " # achieved\n", " state.info['command'] = jp.where(\n", " (done == 1.0) \u0026 (state.info['step'] \u003e self._reset_horizon),\n", " self.sample_command(cmd_rng), state.info['command'])\n", " # reset the step counter when done\n", " state.info['step'] = jp.where(\n", " (done == 1.0) | (state.info['step'] \u003e self._reset_horizon), 0,\n", " state.info['step']\n", " )\n", "\n", " # log total displacement as a proxy metric\n", " state.metrics['total_dist'] = math.normalize(x.pos[self.torso_idx])[1]\n", " for k in state.info['reward_tuple'].keys():\n", " state.metrics[k] = state.info['reward_tuple'][k]\n", "\n", " state = state.replace(\n", " pipeline_state=data, obs=obs + obs_noise, reward=reward,\n", " done=done)\n", " return state\n", "\n", " def _get_obs(self, qpos: jax.Array, x: Transform, xd: Motion,\n", " state_info: Dict[str, Any]) -\u003e jax.Array:\n", " # Get observations:\n", " # yaw_rate, projected_gravity, command, motor_angles, last_action\n", "\n", " inv_base_orientation = math.quat_inv(x.rot[0])\n", " local_rpyrate = math.rotate(xd.ang[0], inv_base_orientation)\n", " cmd = state_info['command']\n", "\n", " obs_list = []\n", " # yaw rate\n", " obs_list.append(jp.array([local_rpyrate[2]]) * 0.25)\n", " # projected gravity\n", " obs_list.append(\n", " math.rotate(jp.array([0.0, 0.0, -1.0]), inv_base_orientation))\n", " # command\n", " obs_list.append(cmd * jp.array([2.0, 2.0, 0.25]))\n", " # motor angles\n", " angles = qpos[7:19]\n", " obs_list.append(angles - self._default_ap_pose)\n", " # last action\n", " obs_list.append(state_info['last_act'])\n", "\n", " obs = jp.clip(jp.concatenate(obs_list), -100.0, 100.0)\n", "\n", " # stack observations through time\n", " single_obs_size = len(obs)\n", " state_info['obs_history'] = jp.roll(\n", " state_info['obs_history'], single_obs_size\n", " )\n", " state_info['obs_history'] = jp.array(\n", " state_info['obs_history']).at[:single_obs_size].set(obs)\n", " return state_info['obs_history']\n", "\n", " # ------------ reward functions----------------\n", " def _reward_lin_vel_z(self, xd: Motion) -\u003e jax.Array:\n", " # Penalize z axis base linear velocity\n", " return jp.square(xd.vel[0, 2])\n", "\n", " def _reward_ang_vel_xy(self, xd: Motion) -\u003e jax.Array:\n", " # Penalize xy axes base angular velocity\n", " return jp.sum(jp.square(xd.ang[0, :2]))\n", "\n", " def _reward_orientation(self, x: Transform) -\u003e jax.Array:\n", " # Penalize non flat base orientation\n", " up = jp.array([0.0, 0.0, 1.0])\n", " rot_up = math.rotate(up, x.rot[0])\n", " return jp.sum(jp.square(rot_up[:2]))\n", "\n", " def _reward_torques(self, torques: jax.Array) -\u003e jax.Array:\n", " # Penalize torques\n", " return jp.sqrt(jp.sum(jp.square(torques))) + jp.sum(jp.abs(torques))\n", "\n", " def _reward_action_rate(\n", " self, act: jax.Array, last_act: jax.Array) -\u003e jax.Array:\n", " # Penalize changes in actions\n", " return jp.sum(jp.square(act - last_act))\n", "\n", " def _reward_tracking_lin_vel(\n", " self, commands: jax.Array, x: Transform, xd: Motion) -\u003e jax.Array:\n", " # Tracking of linear velocity commands (xy axes)\n", " local_vel = math.rotate(xd.vel[0], math.quat_inv(x.rot[0]))\n", " lin_vel_error = jp.sum(jp.square(commands[:2] - local_vel[:2]))\n", " lin_vel_reward = jp.exp(\n", " -lin_vel_error / self.reward_config.rewards.tracking_sigma\n", " )\n", " return lin_vel_reward\n", "\n", " def _reward_tracking_ang_vel(\n", " self, commands: jax.Array, x: Transform, xd: Motion) -\u003e jax.Array:\n", " # Tracking of angular velocity commands (yaw)\n", " base_ang_vel = math.rotate(xd.ang[0], math.quat_inv(x.rot[0]))\n", " ang_vel_error = jp.square(commands[2] - base_ang_vel[2])\n", " return jp.exp(-ang_vel_error/self.reward_config.rewards.tracking_sigma)\n", "\n", " def _reward_feet_air_time(\n", " self, air_time: jax.Array, first_contact: jax.Array,\n", " commands: jax.Array) -\u003e jax.Array:\n", " # Reward air time.\n", " rew_air_time = jp.sum((air_time - 0.1) * first_contact)\n", " rew_air_time *= (\n", " math.normalize(commands[:2])[1] \u003e 0.05\n", " ) # no reward for zero command\n", " return rew_air_time\n", "\n", " def _reward_stand_still(\n", " self, commands: jax.Array, joint_angles: jax.Array,\n", " default_angles: jax.Array) -\u003e jax.Array:\n", " # Penalize motion at zero commands\n", " return jp.sum(jp.abs(joint_angles - default_angles)) * (\n", " math.normalize(commands[:2])[1] \u003c 0.1\n", " )\n", "\n", " def _get_feet_pos_vel(\n", " self, x: Transform, xd: Motion) -\u003e Tuple[jax.Array, jax.Array]:\n", " offset = Transform.create(pos=self._feet_pos)\n", " pos = x.take(self._feet_index).vmap().do(offset).pos\n", " vel = offset.vmap().do(xd.take(self._feet_index)).vel\n", " return pos, vel\n", "\n", " def _reward_foot_slip(\n", " self, x: Transform, xd: Motion, contact_filt: jax.Array) -\u003e jax.Array:\n", " # Get feet velocities\n", " _, foot_world_vel = self._get_feet_pos_vel(x, xd)\n", " # Penalize large feet velocity for feet that are in contact with the ground.\n", " return jp.sum(\n", " jp.square(foot_world_vel[:, :2]) * contact_filt.reshape((-1, 1))\n", " )\n", "\n", "\n", "envs.register_environment('barkour', BarkourEnv)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "pi_yrcz-Qp3W" }, "outputs": [], "source": [ "env_name = 'barkour'\n", "env = envs.get_environment(env_name)\n", "\n", "# re-instantiate the renderer\n", "renderer = mujoco.Renderer(env.model)" ] }, { "cell_type": "markdown", "metadata": { "id": "nxaNFP9mA23H" }, "source": [ "## Train Policy\n", "\n", "To train a policy with domain randomization, we pass in the domain randomization function into the brax train function; brax will call the domain randomization function when rolling out episodes. Training the quadruped takes about 14 minutes on a Tesla V100 GPU." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "cHJCbESGA7Rk" }, "outputs": [], "source": [ "make_networks_factory = functools.partial(\n", " ppo_networks.make_ppo_networks,\n", " policy_hidden_layer_sizes=(128, 128, 128, 128))\n", "train_fn = functools.partial(\n", " ppo.train,\n", " num_timesteps=60_000_000, num_evals=3, reward_scaling=1,\n", " episode_length=1000, normalize_observations=True,\n", " action_repeat=1, unroll_length=20, num_minibatches=8, gae_lambda=0.95,\n", " num_updates_per_batch=4, discounting=0.99, learning_rate=3e-4,\n", " entropy_cost=1e-2, num_envs=8192, batch_size=1024,\n", " network_factory=make_networks_factory,\n", " num_resets_per_eval=10,\n", " randomization_fn=domain_randomize, seed=0)\n", "\n", "\n", "x_data = []\n", "y_data = []\n", "ydataerr = []\n", "times = [datetime.now()]\n", "max_y, min_y = 30, 0\n", "\n", "# Reset environments since internals may be overwritten by tracers from the\n", "# domain randomization function.\n", "env = envs.get_environment(env_name)\n", "eval_env = envs.get_environment(env_name)\n", "make_inference_fn, params, _= train_fn(environment=env,\n", " progress_fn=progress,\n", " eval_env=eval_env)\n", "\n", "print(f'time to jit: {times[1] - times[0]}')\n", "print(f'time to train: {times[-1] - times[1]}')" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "8Yge-CGP5JoO" }, "outputs": [], "source": [ "# Save and reload params.\n", "model_path = '/tmp/mjx_brax_quadruped_policy'\n", "model.save_params(model_path, params)\n", "params = model.load_params(model_path)\n", "\n", "inference_fn = make_inference_fn(params)\n", "jit_inference_fn = jax.jit(inference_fn)" ] }, { "cell_type": "markdown", "metadata": { "id": "L01IrN4oCIkC" }, "source": [ "## Visualize Policy\n", "\n", "For the Barkour Quadruped, the joystick commands can be set through `x_vel`, `y_vel`, and `ang_vel`. `x_vel` and `y_vel` define the linear forward and sideways velocities with respect to the quadruped torso. `ang_vel` defines the angular velocity of the torso in the z direction." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "VTbpEtXnEecd" }, "outputs": [], "source": [ "eval_env = envs.get_environment(env_name)\n", "\n", "jit_reset = jax.jit(eval_env.reset)\n", "jit_step = jax.jit(eval_env.step)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "HRRN-8L-BivZ" }, "outputs": [], "source": [ "\n", "# @markdown Commands **only used for Barkour Env**:\n", "x_vel = 1.0 #@param {type: \"number\"}\n", "y_vel = 0.0 #@param {type: \"number\"}\n", "ang_vel = -0.5 #@param {type: \"number\"}\n", "\n", "the_command = jp.array([x_vel, y_vel, ang_vel])\n", "\n", "# initialize the state\n", "rng = jax.random.PRNGKey(0)\n", "state = jit_reset(rng)\n", "state.info['command'] = the_command\n", "rollout = [state]\n", "images = [get_image(state, camera='track')]\n", "\n", "# grab a trajectory\n", "n_steps = 500\n", "render_every = 2\n", "\n", "for i in range(n_steps):\n", " act_rng, rng = jax.random.split(rng)\n", " ctrl, _ = jit_inference_fn(state.obs, act_rng)\n", " state = jit_step(state, ctrl)\n", " rollout.append(state)\n", " if i % render_every == 0:\n", " images.append(get_image(state, camera='track'))\n", "\n", "media.show_video(images, fps=1.0 / eval_env.dt / render_every)" ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuClass": "premium", "gpuType": "V100", "private_outputs": true, "provenance": [ { "file_id": "1brcF4_qCRS2ASc-QQw1rsEwl5IjzGvq2", "timestamp": 1697763780236 } ], "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }