Google DeepMind
06670155c4
PiperOrigin-RevId: 571066455 Change-Id: I8b0638181cdc751cf9fbacfb84c1ac551d7ce4bd
1272 lines
52 KiB
Plaintext
1272 lines
52 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "MpkYHwCqk7W-"
|
|
},
|
|
"source": [
|
|
"\n",
|
|
"\n",
|
|
"# \u003ch1\u003e\u003ccenter\u003eTutorial \u003ca href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial_mjx_training.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 MJX**](https://github.com/google-deepmind/mujoco#readme), a jax-based version 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 training 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 etc. Notably `State.pipeline_state` holds a `mjx.Data` object, which is analogous to `mjData` in MuJoCo.\n",
|
|
"\n",
|
|
"`MjxEnv` is an implementation of the `brax.envs.base.Env` class that initializes `mjx.Model` and `mjx.Data` objects. Inheriting from `brax.envs.base.Env` allows us to use the training agents implemented in brax. Notice that `MjxEnv` calls `mjx.step` for every `pipeline_step`, which is analgous to `mujoco.mj_step`."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "7DQ_rW4CkIB_"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"#@title State and MjxEnv\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",
|
|
"\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.xipos[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": [
|
|
"Now we can define environment implementations for Humanoid and the [Barkour v0 Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_v0). The environments define the reward and stepping logic. `reset` initializes a `State`, and `step` steps through the physics step and reward logic."
|
|
]
|
|
},
|
|
{
|
|
"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",
|
|
"\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": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"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",
|
|
" # Penalize the change in the action and encourage smooth\n",
|
|
" # actions. L2 regularization |action - last_action|^2\n",
|
|
" action_rate=-0.3,\n",
|
|
" # Encourage long swing steps. However, it does not\n",
|
|
" # encourage high clearances.\n",
|
|
" feet_air_time=0.1,\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",
|
|
"\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.5] # 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": "markdown",
|
|
"metadata": {
|
|
"id": "P1K6IznI2y83"
|
|
},
|
|
"source": [
|
|
"## Visualize a Rollout\n",
|
|
"\n",
|
|
"Let's visualize an environment rollout by running a few env steps with sinusoidal actuation."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "EhKLFK54C1CH"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# instantiate the environment\n",
|
|
"\n",
|
|
"env_name = 'barkour' # @param ['barkour', '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)\n",
|
|
"\n",
|
|
"def get_image(state: 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",
|
|
" if env_name == 'barkour':\n",
|
|
" camera='track'\n",
|
|
" elif env_name == 'humanoid':\n",
|
|
" camera = 'side'\n",
|
|
" else:\n",
|
|
" raise NotImplementedError(env_name)\n",
|
|
" # use the mjData object to update the renderer\n",
|
|
" renderer.update_scene(d, camera=camera)\n",
|
|
" return renderer.render()"
|
|
]
|
|
},
|
|
{
|
|
"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)]\n",
|
|
"\n",
|
|
"# grab a trajectory\n",
|
|
"n_steps = 500\n",
|
|
"render_every = 5\n",
|
|
"\n",
|
|
"for i in range(n_steps):\n",
|
|
" ctrl = jp.ones(env.sys.nu) * jp.sin(2.0 * jp.pi * i / 500)\n",
|
|
" state = jit_step(state, ctrl)\n",
|
|
" rollout.append(state)\n",
|
|
" if i % render_every == 0:\n",
|
|
" images.append(get_image(state))\n",
|
|
"\n",
|
|
"media.show_video(images, fps=1.0 / env.dt / render_every)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "65mIPj6DQNNa"
|
|
},
|
|
"source": [
|
|
"# Domain Randomization\n",
|
|
"\n",
|
|
"We have all the pieces to train a policy, but we might also want to include randomization over certain `mjModel` parameters. 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.8, maxval=1.2)\n",
|
|
" friction = sys.geom_friction.at[:, 0].set(friction)\n",
|
|
" # actuator\n",
|
|
" _, key = jax.random.split(key, 2)\n",
|
|
" gain_range = (-5, 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`."
|
|
]
|
|
},
|
|
{
|
|
"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": "LdqPYBlPVWwc"
|
|
},
|
|
"source": [
|
|
"Note that we'll pass the `domain_randomize` function into the brax trainer; the brax training algorithms will call the randomizer over the environment batch dimension."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "nxaNFP9mA23H"
|
|
},
|
|
"source": [
|
|
"# Train Policy\n",
|
|
"\n",
|
|
"Let's finally train a policy with PPO. Training takes about 12-13 minutes on a Tesla V100 GPU for both environments."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "cHJCbESGA7Rk"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"\n",
|
|
"if env_name == 'barkour':\n",
|
|
" 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=32, gae_lambda=0.95,\n",
|
|
" num_updates_per_batch=4, discounting=0.99, learning_rate=3.0e-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",
|
|
"elif env_name == 'humanoid':\n",
|
|
" 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",
|
|
"else:\n",
|
|
" raise NotImplementedError(\n",
|
|
" f'env_name: {env_name} is not implemented in this notebook.')\n",
|
|
"\n",
|
|
"\n",
|
|
"max_y = {'barkour':30, 'humanoid': 13000}[env_name]\n",
|
|
"min_y = {'barkour': -15}.get(env_name, 0)\n",
|
|
"\n",
|
|
"x_data = []\n",
|
|
"y_data = []\n",
|
|
"ydataerr = []\n",
|
|
"times = [datetime.now()]\n",
|
|
"\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",
|
|
"# Reset environments since internals may be overwritten by tracers due to\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": "markdown",
|
|
"metadata": {
|
|
"id": "-dRmWHsHB_1K"
|
|
},
|
|
"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": "u9sdN4Xa5JoH"
|
|
},
|
|
"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": "8Yge-CGP5JoO"
|
|
},
|
|
"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": "L01IrN4oCIkC"
|
|
},
|
|
"source": [
|
|
"# Visualize Policy\n",
|
|
"\n",
|
|
"Finally we can visualize the policy. 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.0 #@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)]\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))\n",
|
|
"\n",
|
|
"media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
|
|
]
|
|
},
|
|
{
|
|
"metadata": {
|
|
"id": "lLXVB0Q3EaC-"
|
|
},
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"We can also perform the physics step using the original MuJoCo python bindings to show that a policy trained in MJX works in MuJoCo-C."
|
|
]
|
|
},
|
|
{
|
|
"metadata": {
|
|
"id": "dXHTu5AIEaC-"
|
|
},
|
|
"cell_type": "code",
|
|
"source": [
|
|
"mj_model = eval_env.model\n",
|
|
"mj_data = mujoco.MjData(mj_model)\n",
|
|
"\n",
|
|
"camera = 'track' if env_name == 'barkour' else 'side'\n",
|
|
"renderer = mujoco.Renderer(mj_model)\n",
|
|
"\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",
|
|
" 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=camera)\n",
|
|
" images.append(renderer.render())\n",
|
|
"\n",
|
|
"media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
|
|
],
|
|
"outputs": [],
|
|
"execution_count": null
|
|
}
|
|
],
|
|
"metadata": {
|
|
"accelerator": "GPU",
|
|
"colab": {
|
|
"gpuClass": "premium",
|
|
"gpuType": "V100",
|
|
"private_outputs": true,
|
|
"provenance": [],
|
|
"toc_visible": true
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"name": "python3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 0
|
|
}
|