Update MJX colab with hfield example.

PiperOrigin-RevId: 641061897
Change-Id: Ia1002173adce7b7b12203857abed18568ae236f1

mjx/tutorial.ipynb

@@ -62,7 +62,7 @@
    "source": [
     "!pip install mujoco\n",
     "!pip install mujoco_mjx\n",
-    "!pip install brax"
+    "!pip install brax\n"
    ]
   },
   {
@@ -102,11 +102,6 @@
     "}\n",
     "\"\"\")\n",
     "\n",
-    "# Tell XLA to use Triton GEMM, this improves steps/sec by ~30% on some GPUs\n",
-    "xla_flags = os.environ.get('XLA_FLAGS', '')\n",
-    "xla_flags += ' --xla_gpu_triton_gemm_any=True'\n",
-    "os.environ['XLA_FLAGS'] = xla_flags\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",
@@ -122,7 +117,12 @@
     "      '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.')"
+    "print('Installation successful.')\n",
+    "\n",
+    "# Tell XLA to use Triton GEMM, this improves steps/sec by ~30% on some GPUs\n",
+    "xla_flags = os.environ.get('XLA_FLAGS', '')\n",
+    "xla_flags += ' --xla_gpu_triton_gemm_any=True'\n",
+    "os.environ['XLA_FLAGS'] = xla_flags\n"
    ]
   },
   {
@@ -155,20 +155,32 @@
    "cell_type": "code",
    "execution_count": 0,
    "metadata": {
+    "cellView": "form",
     "id": "ObF1UXrkb0Nd"
    },
    "outputs": [],
    "source": [
     "#@title Import MuJoCo, MJX, and Brax\n",
-    "\n",
-    "\n",
     "from datetime import datetime\n",
+    "from etils import epath\n",
     "import functools\n",
     "from IPython.display import HTML\n",
+    "from typing import Any, Dict, Sequence, Tuple, Union\n",
+    "import os\n",
+    "from ml_collections import config_dict\n",
+    "\n",
+    "\n",
     "import jax\n",
     "from jax import numpy as jp\n",
     "import numpy as np\n",
-    "from typing import Any, Dict, Sequence, Tuple, Union\n",
+    "from flax.training import orbax_utils\n",
+    "from flax import struct\n",
+    "from matplotlib import pyplot as plt\n",
+    "import mediapy as media\n",
+    "from orbax import checkpoint as ocp\n",
+    "\n",
+    "import mujoco\n",
+    "from mujoco import mjx\n",
     "\n",
     "from brax import base\n",
     "from brax import envs\n",
@@ -178,15 +190,7 @@
     "from brax.mjx.base import State as MjxState\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 html, mjcf, model\n",
-    "\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"
+    "from brax.io import html, mjcf, model\n"
    ]
   },
   {
@@ -892,7 +896,7 @@
    },
    "outputs": [],
    "source": [
-    "!git clone https://github.com/google-deepmind/mujoco_menagerie"
+    "!git clone https://github.com/google-deepmind/mujoco_menagerie\n"
    ]
   },
   {
@@ -972,9 +976,11 @@
     "      obs_noise: float = 0.05,\n",
     "      action_scale: float = 0.3,\n",
     "      kick_vel: float = 0.05,\n",
+    "      scene_file: str = 'scene_mjx.xml',\n",
     "      **kwargs,\n",
     "  ):\n",
-    "    path = epath.Path('mujoco_menagerie/google_barkour_vb/scene_mjx.xml')\n",
+    "# path = epath.Path('mujoco_menagerie/google_barkour_vb')\n",
+    "    path = path / scene_file\n",
     "    sys = mjcf.load(path.as_posix())\n",
     "    self._dt = 0.02  # this environment is 50 fps\n",
     "    sys = sys.tree_replace({'opt.timestep': 0.004})\n",
@@ -1284,10 +1290,11 @@
     "    return done & (step < 500)\n",
     "\n",
     "  def render(\n",
-    "      self, trajectory: List[base.State], camera: str | None = None\n",
+    "      self, trajectory: List[base.State], camera: str | None = None,\n",
+    "      width: int = 240, height: int = 320,\n",
     "  ) -> Sequence[np.ndarray]:\n",
     "    camera = camera or 'track'\n",
-    "    return super().render(trajectory, camera=camera)\n",
+    "    return super().render(trajectory, camera=camera, width=width, height=height)\n",
     "\n",
     "envs.register_environment('barkour', BarkourEnv)"
    ]
@@ -1323,6 +1330,17 @@
    },
    "outputs": [],
    "source": [
+    "ckpt_path = epath.Path('/tmp/quadrupred_joystick/ckpts')\n",
+    "ckpt_path.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "def policy_params_fn(current_step, make_policy, params):\n",
+    "  # save checkpoints\n",
+    "  orbax_checkpointer = ocp.PyTreeCheckpointer()\n",
+    "  save_args = orbax_utils.save_args_from_target(params)\n",
+    "  path = ckpt_path / f'{current_step}'\n",
+    "  orbax_checkpointer.save(path, params, force=True, save_args=save_args)\n",
+    "\n",
+    "\n",
     "make_networks_factory = functools.partial(\n",
     "    ppo_networks.make_ppo_networks,\n",
     "        policy_hidden_layer_sizes=(128, 128, 128, 128))\n",
@@ -1333,7 +1351,9 @@
     "      num_updates_per_batch=4, discounting=0.97, learning_rate=3.0e-4,\n",
     "      entropy_cost=1e-2, num_envs=8192, batch_size=256,\n",
     "      network_factory=make_networks_factory,\n",
-    "      randomization_fn=domain_randomize, seed=0)\n",
+    "      randomization_fn=domain_randomize,\n",
+    "      policy_params_fn=policy_params_fn,\n",
+    "      seed=0)\n",
     "\n",
     "x_data = []\n",
     "y_data = []\n",
@@ -1450,6 +1470,136 @@
    "source": [
     "HTML(html.render(eval_env.sys.tree_replace({'opt.timestep': eval_env.dt}), rollout))"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "gNagGnBODotY"
+   },
+   "source": [
+    "## Train Policy with Height Field\n",
+    "\n",
+    "We may also want the quadruped to learn to walk on rought terrain. Let's take the latest checkpoint from the joystick policy above, and finetune it on a height field terrain."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "wlT3xLouKxqT"
+   },
+   "outputs": [],
+   "source": [
+    "# use the height field scene\n",
+    "scene_file = 'scene_hfield_mjx.xml'\n",
+    "\n",
+    "env = envs.get_environment(env_name, scene_file=scene_file)\n",
+    "jit_reset = jax.jit(env.reset)\n",
+    "state = jit_reset(jax.random.PRNGKey(0))\n",
+    "plt.imshow(env.render([state.pipeline_state], camera='track')[0])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "jOTx-OyPDqPW"
+   },
+   "outputs": [],
+   "source": [
+    "# grab the latest checkpoint from the flat terrain joystick policy\n",
+    "latest_ckpts = list(ckpt_path.glob('*'))\n",
+    "latest_ckpts.sort()\n",
+    "latest_ckpt = latest_ckpts[0]\n",
+    "\n",
+    "train_fn = functools.partial(\n",
+    "      ppo.train, num_timesteps=40_000_000, num_evals=5,\n",
+    "      reward_scaling=1, episode_length=1000, normalize_observations=True,\n",
+    "      action_repeat=1, unroll_length=20, num_minibatches=32,\n",
+    "      num_updates_per_batch=4, discounting=0.97, learning_rate=3.0e-4,\n",
+    "      entropy_cost=1e-2, num_envs=8192, batch_size=256,\n",
+    "      network_factory=make_networks_factory,\n",
+    "      randomization_fn=domain_randomize, seed=0,\n",
+    "      restore_checkpoint_path=latest_ckpt)\n",
+    "\n",
+    "x_data = []\n",
+    "y_data = []\n",
+    "ydataerr = []\n",
+    "times = [datetime.now()]\n",
+    "max_y, min_y = 40, 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, scene_file=scene_file)\n",
+    "eval_env = envs.get_environment(env_name, scene_file=scene_file)\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": "0wHRJyjLFww6"
+   },
+   "source": [
+    "## Visualize Policy with Height Field"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "1X57XkaVFu-v"
+   },
+   "outputs": [],
+   "source": [
+    "eval_env = envs.get_environment(env_name, scene_file=scene_file)\n",
+    "\n",
+    "jit_reset = jax.jit(eval_env.reset)\n",
+    "jit_step = jax.jit(eval_env.step)\n",
+    "inference_fn = make_inference_fn(params)\n",
+    "jit_inference_fn = jax.jit(inference_fn)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "id": "nAxexZcVFu-v"
+   },
+   "outputs": [],
+   "source": [
+    "# @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.pipeline_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.pipeline_state)\n",
+    "\n",
+    "media.show_video(\n",
+    "    eval_env.render(rollout[::render_every], camera='track'),\n",
+    "    fps=1.0 / eval_env.dt / render_every)"
+   ]
   }
  ],
  "metadata": {