Reformat notebook and add back sim2sim.
PiperOrigin-RevId: 575246433 Change-Id: I8627916e7dedd5cf64dab6e0151f41c3153e7a68
This commit is contained in:
Baruch Tabanpour
committed by
Copybara-Service
Copybara-Service
parent
ebe081f27e
commit
e57b34303a
+433
-246
@@ -187,11 +187,9 @@
|
||||
},
|
||||
"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",
|
||||
"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 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`."
|
||||
"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"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -202,7 +200,7 @@
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title State and MjxEnv\n",
|
||||
"#@title State\n",
|
||||
"\n",
|
||||
"@struct.dataclass\n",
|
||||
"class State(Base):\n",
|
||||
@@ -223,8 +221,30 @@
|
||||
" 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",
|
||||
" 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",
|
||||
@@ -305,7 +325,7 @@
|
||||
"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."
|
||||
"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."
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -444,6 +464,372 @@
|
||||
"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,
|
||||
@@ -483,13 +869,14 @@
|
||||
" # 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.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.3,\n",
|
||||
" action_rate=-0.1,\n",
|
||||
" # Encourage long swing steps. However, it does not\n",
|
||||
" # encourage high clearances.\n",
|
||||
" feet_air_time=0.1,\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",
|
||||
@@ -550,7 +937,7 @@
|
||||
" 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_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",
|
||||
@@ -857,171 +1244,19 @@
|
||||
"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"
|
||||
"id": "pi_yrcz-Qp3W"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# instantiate the environment\n",
|
||||
"\n",
|
||||
"env_name = 'humanoid' # @param ['barkour', 'humanoid']\n",
|
||||
"env_name = 'barkour'\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."
|
||||
"# re-instantiate the renderer\n",
|
||||
"renderer = mujoco.Renderer(env.model)"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -1030,9 +1265,9 @@
|
||||
"id": "nxaNFP9mA23H"
|
||||
},
|
||||
"source": [
|
||||
"# Train Policy\n",
|
||||
"## 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."
|
||||
"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."
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -1043,59 +1278,28 @@
|
||||
},
|
||||
"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",
|
||||
"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",
|
||||
"max_y, min_y = 30, 0\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",
|
||||
"# 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",
|
||||
@@ -1107,30 +1311,6 @@
|
||||
"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,
|
||||
@@ -1139,7 +1319,9 @@
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title Load Model and Define Inference Function\n",
|
||||
"# 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",
|
||||
@@ -1152,9 +1334,9 @@
|
||||
"id": "L01IrN4oCIkC"
|
||||
},
|
||||
"source": [
|
||||
"# Visualize Policy\n",
|
||||
"## 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."
|
||||
"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."
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -1183,7 +1365,7 @@
|
||||
"# @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",
|
||||
"ang_vel = -0.5 #@param {type: \"number\"}\n",
|
||||
"\n",
|
||||
"the_command = jp.array([x_vel, y_vel, ang_vel])\n",
|
||||
"\n",
|
||||
@@ -1192,7 +1374,7 @@
|
||||
"state = jit_reset(rng)\n",
|
||||
"state.info['command'] = the_command\n",
|
||||
"rollout = [state]\n",
|
||||
"images = [get_image(state)]\n",
|
||||
"images = [get_image(state, camera='track')]\n",
|
||||
"\n",
|
||||
"# grab a trajectory\n",
|
||||
"n_steps = 500\n",
|
||||
@@ -1204,7 +1386,7 @@
|
||||
" state = jit_step(state, ctrl)\n",
|
||||
" rollout.append(state)\n",
|
||||
" if i % render_every == 0:\n",
|
||||
" images.append(get_image(state))\n",
|
||||
" images.append(get_image(state, camera='track'))\n",
|
||||
"\n",
|
||||
"media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
|
||||
]
|
||||
@@ -1216,7 +1398,12 @@
|
||||
"gpuClass": "premium",
|
||||
"gpuType": "V100",
|
||||
"private_outputs": true,
|
||||
"provenance": [],
|
||||
"provenance": [
|
||||
{
|
||||
"file_id": "1brcF4_qCRS2ASc-QQw1rsEwl5IjzGvq2",
|
||||
"timestamp": 1697763780236
|
||||
}
|
||||
],
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
|
||||
Reference in New Issue
Block a user