chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Tidy up Colab notebook files.

Browse Source 
PiperOrigin-RevId: 613163325
Change-Id: Iec56375bb7631b8a3035f884b3aad6b54fb7b88b
This commit is contained in:
Nimrod Gileadi
2024-03-06 04:13:02 -08:00
committed by Copybara-Service
parent e2c14188ad
commit a00f1b5f30
3 changed files with 4598 additions and 4610 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/tutorial.ipynb
+85 -91
View File
@@ -8,11 +8,11 @@
"source": [ "source": [
"![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n", "![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n",
"\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", "# <h1><center>Tutorial <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a></center></h1>\n",
"\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", "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", "\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", "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime > Change runtime type\".\n",
"\n", "\n",
"\n", "\n",
"\n", "\n",
@@ -38,9 +38,9 @@
"id": "_UbO9uhtBSX5" "id": "_UbO9uhtBSX5"
}, },
"source": [ "source": [
"\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eCopyright 2023 DeepMind Technologies Limited.\u003c/small\u003e\u003c/p\u003e\n", "> <p><small><small>Copyright 2023 DeepMind Technologies Limited.</small></p>\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", "> <p><small><small>Licensed 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 <a href=\"http://www.apache.org/licenses/LICENSE-2.0\">http://www.apache.org/licenses/LICENSE-2.0</a>.</small></small></p>\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" "> <p><small><small>Unless 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.</small></small></p>"
] ]
}, },
{ {
@@ -54,7 +54,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "Xqo7pyX-n72M" "id": "Xqo7pyX-n72M"
}, },
@@ -67,7 +67,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"cellView": "form", "cellView": "form",
"id": "IbZxYDxzoz5R" "id": "IbZxYDxzoz5R"
@@ -114,7 +114,7 @@
"try:\n", "try:\n",
" print('Checking that the installation succeeded:')\n", " print('Checking that the installation succeeded:')\n",
" import mujoco\n", " import mujoco\n",
" mujoco.MjModel.from_xml_string('\u003cmujoco/\u003e')\n", " mujoco.MjModel.from_xml_string('<mujoco/>')\n",
"except Exception as e:\n", "except Exception as e:\n",
" raise e from RuntimeError(\n", " raise e from RuntimeError(\n",
" 'Something went wrong during installation. Check the shell output above '\n", " 'Something went wrong during installation. Check the shell output above '\n",
@@ -127,7 +127,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"cellView": "form", "cellView": "form",
"id": "T5f4w3Kq2X14" "id": "T5f4w3Kq2X14"
@@ -142,7 +142,7 @@
"\n", "\n",
"# Graphics and plotting.\n", "# Graphics and plotting.\n",
"print('Installing mediapy:')\n", "print('Installing mediapy:')\n",
"!command -v ffmpeg \u003e/dev/null || (apt update \u0026\u0026 apt install -y ffmpeg)\n", "!command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
"!pip install -q mediapy\n", "!pip install -q mediapy\n",
"import mediapy as media\n", "import mediapy as media\n",
"import matplotlib.pyplot as plt\n", "import matplotlib.pyplot as plt\n",
@@ -153,7 +153,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "ObF1UXrkb0Nd" "id": "ObF1UXrkb0Nd"
}, },
@@ -204,23 +204,23 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "bNus3mbbDz6a" "id": "bNus3mbbDz6a"
}, },
"outputs": [], "outputs": [],
"source": [ "source": [
"xml = \"\"\"\n", "xml = \"\"\"\n",
"\u003cmujoco\u003e\n", "<mujoco>\n",
" \u003cworldbody\u003e\n", " <worldbody>\n",
" \u003clight name=\"top\" pos=\"0 0 1\"/\u003e\n", " <light name=\"top\" pos=\"0 0 1\"/>\n",
" \u003cbody name=\"box_and_sphere\" euler=\"0 0 -30\"\u003e\n", " <body name=\"box_and_sphere\" euler=\"0 0 -30\">\n",
" \u003cjoint name=\"swing\" type=\"hinge\" axis=\"1 -1 0\" pos=\"-.2 -.2 -.2\"/\u003e\n", " <joint name=\"swing\" type=\"hinge\" axis=\"1 -1 0\" pos=\"-.2 -.2 -.2\"/>\n",
" \u003cgeom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/\u003e\n", " <geom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/>\n",
" \u003cgeom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/\u003e\n", " <geom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/>\n",
" \u003c/body\u003e\n", " </body>\n",
" \u003c/worldbody\u003e\n", " </worldbody>\n",
"\u003c/mujoco\u003e\n", "</mujoco>\n",
"\"\"\"\n", "\"\"\"\n",
"\n", "\n",
"# Make model, data, and renderer\n", "# Make model, data, and renderer\n",
@@ -240,7 +240,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "TSpoOWqeEC3P" "id": "TSpoOWqeEC3P"
}, },
@@ -261,7 +261,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "ZOD582pfOLP-" "id": "ZOD582pfOLP-"
}, },
@@ -282,7 +282,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "HDlPlX05I3m-" "id": "HDlPlX05I3m-"
}, },
@@ -297,9 +297,9 @@
"\n", "\n",
"frames = []\n", "frames = []\n",
"mujoco.mj_resetData(mj_model, mj_data)\n", "mujoco.mj_resetData(mj_model, mj_data)\n",
"while mj_data.time \u003c duration:\n", "while mj_data.time < duration:\n",
" mujoco.mj_step(mj_model, mj_data)\n", " mujoco.mj_step(mj_model, mj_data)\n",
" if len(frames) \u003c mj_data.time * framerate:\n", " if len(frames) < mj_data.time * framerate:\n",
" renderer.update_scene(mj_data, scene_option=scene_option)\n", " renderer.update_scene(mj_data, scene_option=scene_option)\n",
" pixels = renderer.render()\n", " pixels = renderer.render()\n",
" frames.append(pixels)\n", " frames.append(pixels)\n",
@@ -321,7 +321,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "Pr29xq0-JRQv" "id": "Pr29xq0-JRQv"
}, },
@@ -333,9 +333,9 @@
"frames = []\n", "frames = []\n",
"mujoco.mj_resetData(mj_model, mj_data)\n", "mujoco.mj_resetData(mj_model, mj_data)\n",
"mjx_data = mjx.put_data(mj_model, mj_data)\n", "mjx_data = mjx.put_data(mj_model, mj_data)\n",
"while mjx_data.time \u003c duration:\n", "while mjx_data.time < duration:\n",
" mjx_data = jit_step(mjx_model, mjx_data)\n", " mjx_data = jit_step(mjx_model, mjx_data)\n",
" if len(frames) \u003c mjx_data.time * framerate:\n", " if len(frames) < mjx_data.time * framerate:\n",
" mj_data = mjx.get_data(mj_model, mjx_data)\n", " mj_data = mjx.get_data(mj_model, mjx_data)\n",
" renderer.update_scene(mj_data, scene_option=scene_option)\n", " renderer.update_scene(mj_data, scene_option=scene_option)\n",
" pixels = renderer.render()\n", " pixels = renderer.render()\n",
@@ -357,7 +357,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "rrdrcKRVK6w9" "id": "rrdrcKRVK6w9"
}, },
@@ -384,7 +384,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "Jtz7j1PDOnw5" "id": "Jtz7j1PDOnw5"
}, },
@@ -409,7 +409,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "mtGMYNLE3QJN" "id": "mtGMYNLE3QJN"
}, },
@@ -458,7 +458,7 @@
" exclude_current_positions_from_observation\n", " exclude_current_positions_from_observation\n",
" )\n", " )\n",
"\n", "\n",
" def reset(self, rng: jp.ndarray) -\u003e State:\n", " def reset(self, rng: jp.ndarray) -> State:\n",
" \"\"\"Resets the environment to an initial state.\"\"\"\n", " \"\"\"Resets the environment to an initial state.\"\"\"\n",
" rng, rng1, rng2 = jax.random.split(rng, 3)\n", " rng, rng1, rng2 = jax.random.split(rng, 3)\n",
"\n", "\n",
@@ -487,7 +487,7 @@
" }\n", " }\n",
" return State(data, obs, reward, done, metrics)\n", " return State(data, obs, reward, done, metrics)\n",
"\n", "\n",
" def step(self, state: State, action: jp.ndarray) -\u003e State:\n", " def step(self, state: State, action: jp.ndarray) -> State:\n",
" \"\"\"Runs one timestep of the environment's dynamics.\"\"\"\n", " \"\"\"Runs one timestep of the environment's dynamics.\"\"\"\n",
" data0 = state.pipeline_state\n", " data0 = state.pipeline_state\n",
" data = self.pipeline_step(data0, action)\n", " data = self.pipeline_step(data0, action)\n",
@@ -498,8 +498,8 @@
" forward_reward = self._forward_reward_weight * velocity[0]\n", " forward_reward = self._forward_reward_weight * velocity[0]\n",
"\n", "\n",
" min_z, max_z = self._healthy_z_range\n", " min_z, max_z = self._healthy_z_range\n",
" is_healthy = jp.where(data.q[2] \u003c min_z, 0.0, 1.0)\n", " is_healthy = jp.where(data.q[2] < min_z, 0.0, 1.0)\n",
" is_healthy = jp.where(data.q[2] \u003e max_z, 0.0, is_healthy)\n", " is_healthy = jp.where(data.q[2] > max_z, 0.0, is_healthy)\n",
" if self._terminate_when_unhealthy:\n", " if self._terminate_when_unhealthy:\n",
" healthy_reward = self._healthy_reward\n", " healthy_reward = self._healthy_reward\n",
" else:\n", " else:\n",
@@ -528,7 +528,7 @@
"\n", "\n",
" def _get_obs(\n", " def _get_obs(\n",
" self, data: mjx.Data, action: jp.ndarray\n", " self, data: mjx.Data, action: jp.ndarray\n",
" ) -\u003e jp.ndarray:\n", " ) -> jp.ndarray:\n",
" \"\"\"Observes humanoid body position, velocities, and angles.\"\"\"\n", " \"\"\"Observes humanoid body position, velocities, and angles.\"\"\"\n",
" position = data.qpos\n", " position = data.qpos\n",
" if self._exclude_current_positions_from_observation:\n", " if self._exclude_current_positions_from_observation:\n",
@@ -562,7 +562,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "EhKLFK54C1CH" "id": "EhKLFK54C1CH"
}, },
@@ -579,7 +579,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "Ph8u-v2Q2xLS" "id": "Ph8u-v2Q2xLS"
}, },
@@ -611,7 +611,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "xLiddQYPApBw" "id": "xLiddQYPApBw"
}, },
@@ -660,14 +660,14 @@
"id": "YYIch0HEApBx" "id": "YYIch0HEApBx"
}, },
"source": [ "source": [
"\u003c!-- ## Save and Load Policy --\u003e\n", "<!-- ## Save and Load Policy -->\n",
"\n", "\n",
"We can save and load the policy using the brax model API." "We can save and load the policy using the brax model API."
] ]
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "Z8gI6qH6ApBx" "id": "Z8gI6qH6ApBx"
}, },
@@ -680,7 +680,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "h4reaWgxApBx" "id": "h4reaWgxApBx"
}, },
@@ -706,7 +706,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "osYasMw4ApBy" "id": "osYasMw4ApBy"
}, },
@@ -720,7 +720,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "d-UhypudApBy" "id": "d-UhypudApBy"
}, },
@@ -760,7 +760,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "w6ixFi4dApBy" "id": "w6ixFi4dApBy"
}, },
@@ -803,7 +803,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "h8mhzKjHQuoL" "id": "h8mhzKjHQuoL"
}, },
@@ -856,7 +856,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "1K45Kp2ASV9s" "id": "1K45Kp2ASV9s"
}, },
@@ -886,7 +886,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "VfyK73gtRXid" "id": "VfyK73gtRXid"
}, },
@@ -897,7 +897,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "y79PoJOCIl-O" "id": "y79PoJOCIl-O"
}, },
@@ -1032,7 +1032,7 @@
" self._foot_radius = 0.0175\n", " self._foot_radius = 0.0175\n",
" self._nv = sys.nv\n", " self._nv = sys.nv\n",
"\n", "\n",
" def sample_command(self, rng: jax.Array) -\u003e jax.Array:\n", " def sample_command(self, rng: jax.Array) -> jax.Array:\n",
" lin_vel_x = [-0.6, 1.5] # min max [m/s]\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", " lin_vel_y = [-0.8, 0.8] # min max [m/s]\n",
" ang_vel_yaw = [-0.7, 0.7] # min max [rad/s]\n", " ang_vel_yaw = [-0.7, 0.7] # min max [rad/s]\n",
@@ -1050,7 +1050,7 @@
" new_cmd = jp.array([lin_vel_x[0], lin_vel_y[0], ang_vel_yaw[0]])\n", " new_cmd = jp.array([lin_vel_x[0], lin_vel_y[0], ang_vel_yaw[0]])\n",
" return new_cmd\n", " return new_cmd\n",
"\n", "\n",
" def reset(self, rng: jax.Array) -\u003e State: # pytype: disable=signature-mismatch\n", " def reset(self, rng: jax.Array) -> State: # pytype: disable=signature-mismatch\n",
" rng, key = jax.random.split(rng)\n", " rng, key = jax.random.split(rng)\n",
"\n", "\n",
" pipeline_state = self.pipeline_init(self._init_q, jp.zeros(self._nv))\n", " pipeline_state = self.pipeline_init(self._init_q, jp.zeros(self._nv))\n",
@@ -1076,7 +1076,7 @@
" state = State(pipeline_state, obs, reward, done, metrics, state_info) # pytype: disable=wrong-arg-types\n", " state = State(pipeline_state, obs, reward, done, metrics, state_info) # pytype: disable=wrong-arg-types\n",
" return state\n", " return state\n",
"\n", "\n",
" def step(self, state: State, action: jax.Array) -\u003e State: # pytype: disable=signature-mismatch\n", " def step(self, state: State, action: jax.Array) -> State: # pytype: disable=signature-mismatch\n",
" rng, cmd_rng, kick_noise_2 = jax.random.split(state.info['rng'], 3)\n", " rng, cmd_rng, kick_noise_2 = jax.random.split(state.info['rng'], 3)\n",
"\n", "\n",
" # kick\n", " # kick\n",
@@ -1102,18 +1102,18 @@
" # foot contact data based on z-position\n", " # foot contact data based on z-position\n",
" foot_pos = pipeline_state.site_xpos[self._feet_site_id] # pytype: disable=attribute-error\n", " foot_pos = pipeline_state.site_xpos[self._feet_site_id] # pytype: disable=attribute-error\n",
" foot_contact_z = foot_pos[:, 2] - self._foot_radius\n", " foot_contact_z = foot_pos[:, 2] - self._foot_radius\n",
" contact = foot_contact_z \u003c 1e-3 # a mm or less off the floor\n", " contact = foot_contact_z < 1e-3 # a mm or less off the floor\n",
" contact_filt_mm = contact | state.info['last_contact']\n", " contact_filt_mm = contact | state.info['last_contact']\n",
" contact_filt_cm = (foot_contact_z \u003c 3e-2) | state.info['last_contact']\n", " contact_filt_cm = (foot_contact_z < 3e-2) | state.info['last_contact']\n",
" first_contact = (state.info['feet_air_time'] \u003e 0) * contact_filt_mm\n", " first_contact = (state.info['feet_air_time'] > 0) * contact_filt_mm\n",
" state.info['feet_air_time'] += self.dt\n", " state.info['feet_air_time'] += self.dt\n",
"\n", "\n",
" # done if joint limits are reached or robot is falling\n", " # done if joint limits are reached or robot is falling\n",
" up = jp.array([0.0, 0.0, 1.0])\n", " up = jp.array([0.0, 0.0, 1.0])\n",
" done = jp.dot(math.rotate(up, x.rot[self._torso_idx - 1]), up) \u003c 0\n", " done = jp.dot(math.rotate(up, x.rot[self._torso_idx - 1]), up) < 0\n",
" done |= jp.any(joint_angles \u003c self.lowers)\n", " done |= jp.any(joint_angles < self.lowers)\n",
" done |= jp.any(joint_angles \u003e self.uppers)\n", " done |= jp.any(joint_angles > self.uppers)\n",
" done |= pipeline_state.x.pos[self._torso_idx - 1, 2] \u003c 0.18\n", " done |= pipeline_state.x.pos[self._torso_idx - 1, 2] < 0.18\n",
"\n", "\n",
" # reward\n", " # reward\n",
" rewards = {\n", " rewards = {\n",
@@ -1156,13 +1156,13 @@
"\n", "\n",
" # sample new command if more than 500 timesteps achieved\n", " # sample new command if more than 500 timesteps achieved\n",
" state.info['command'] = jp.where(\n", " state.info['command'] = jp.where(\n",
" state.info['step'] \u003e 500,\n", " state.info['step'] > 500,\n",
" self.sample_command(cmd_rng),\n", " self.sample_command(cmd_rng),\n",
" state.info['command'],\n", " state.info['command'],\n",
" )\n", " )\n",
" # reset the step counter when done\n", " # reset the step counter when done\n",
" state.info['step'] = jp.where(\n", " state.info['step'] = jp.where(\n",
" done | (state.info['step'] \u003e 500), 0, state.info['step']\n", " done | (state.info['step'] > 500), 0, state.info['step']\n",
" )\n", " )\n",
"\n", "\n",
" # log total displacement as a proxy metric\n", " # log total displacement as a proxy metric\n",
@@ -1180,7 +1180,7 @@
" pipeline_state: base.State,\n", " pipeline_state: base.State,\n",
" state_info: dict[str, Any],\n", " state_info: dict[str, Any],\n",
" obs_history: jax.Array,\n", " obs_history: jax.Array,\n",
" ) -\u003e jax.Array:\n", " ) -> jax.Array:\n",
" inv_torso_rot = math.quat_inv(pipeline_state.x.rot[0])\n", " inv_torso_rot = math.quat_inv(pipeline_state.x.rot[0])\n",
" local_rpyrate = math.rotate(pipeline_state.xd.ang[0], inv_torso_rot)\n", " local_rpyrate = math.rotate(pipeline_state.xd.ang[0], inv_torso_rot)\n",
"\n", "\n",
@@ -1202,33 +1202,33 @@
" return obs\n", " return obs\n",
"\n", "\n",
" # ------------ reward functions----------------\n", " # ------------ reward functions----------------\n",
" def _reward_lin_vel_z(self, xd: Motion) -\u003e jax.Array:\n", " def _reward_lin_vel_z(self, xd: Motion) -> jax.Array:\n",
" # Penalize z axis base linear velocity\n", " # Penalize z axis base linear velocity\n",
" return jp.square(xd.vel[0, 2])\n", " return jp.square(xd.vel[0, 2])\n",
"\n", "\n",
" def _reward_ang_vel_xy(self, xd: Motion) -\u003e jax.Array:\n", " def _reward_ang_vel_xy(self, xd: Motion) -> jax.Array:\n",
" # Penalize xy axes base angular velocity\n", " # Penalize xy axes base angular velocity\n",
" return jp.sum(jp.square(xd.ang[0, :2]))\n", " return jp.sum(jp.square(xd.ang[0, :2]))\n",
"\n", "\n",
" def _reward_orientation(self, x: Transform) -\u003e jax.Array:\n", " def _reward_orientation(self, x: Transform) -> jax.Array:\n",
" # Penalize non flat base orientation\n", " # Penalize non flat base orientation\n",
" up = jp.array([0.0, 0.0, 1.0])\n", " up = jp.array([0.0, 0.0, 1.0])\n",
" rot_up = math.rotate(up, x.rot[0])\n", " rot_up = math.rotate(up, x.rot[0])\n",
" return jp.sum(jp.square(rot_up[:2]))\n", " return jp.sum(jp.square(rot_up[:2]))\n",
"\n", "\n",
" def _reward_torques(self, torques: jax.Array) -\u003e jax.Array:\n", " def _reward_torques(self, torques: jax.Array) -> jax.Array:\n",
" # Penalize torques\n", " # Penalize torques\n",
" return jp.sqrt(jp.sum(jp.square(torques))) + jp.sum(jp.abs(torques))\n", " return jp.sqrt(jp.sum(jp.square(torques))) + jp.sum(jp.abs(torques))\n",
"\n", "\n",
" def _reward_action_rate(\n", " def _reward_action_rate(\n",
" self, act: jax.Array, last_act: jax.Array\n", " self, act: jax.Array, last_act: jax.Array\n",
" ) -\u003e jax.Array:\n", " ) -> jax.Array:\n",
" # Penalize changes in actions\n", " # Penalize changes in actions\n",
" return jp.sum(jp.square(act - last_act))\n", " return jp.sum(jp.square(act - last_act))\n",
"\n", "\n",
" def _reward_tracking_lin_vel(\n", " def _reward_tracking_lin_vel(\n",
" self, commands: jax.Array, x: Transform, xd: Motion\n", " self, commands: jax.Array, x: Transform, xd: Motion\n",
" ) -\u003e jax.Array:\n", " ) -> jax.Array:\n",
" # Tracking of linear velocity commands (xy axes)\n", " # Tracking of linear velocity commands (xy axes)\n",
" local_vel = math.rotate(xd.vel[0], math.quat_inv(x.rot[0]))\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_error = jp.sum(jp.square(commands[:2] - local_vel[:2]))\n",
@@ -1239,7 +1239,7 @@
"\n", "\n",
" def _reward_tracking_ang_vel(\n", " def _reward_tracking_ang_vel(\n",
" self, commands: jax.Array, x: Transform, xd: Motion\n", " self, commands: jax.Array, x: Transform, xd: Motion\n",
" ) -\u003e jax.Array:\n", " ) -> jax.Array:\n",
" # Tracking of angular velocity commands (yaw)\n", " # Tracking of angular velocity commands (yaw)\n",
" base_ang_vel = math.rotate(xd.ang[0], math.quat_inv(x.rot[0]))\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", " ang_vel_error = jp.square(commands[2] - base_ang_vel[2])\n",
@@ -1247,11 +1247,11 @@
"\n", "\n",
" def _reward_feet_air_time(\n", " def _reward_feet_air_time(\n",
" self, air_time: jax.Array, first_contact: jax.Array, commands: jax.Array\n", " self, air_time: jax.Array, first_contact: jax.Array, commands: jax.Array\n",
" ) -\u003e jax.Array:\n", " ) -> jax.Array:\n",
" # Reward air time.\n", " # Reward air time.\n",
" rew_air_time = jp.sum((air_time - 0.1) * first_contact)\n", " rew_air_time = jp.sum((air_time - 0.1) * first_contact)\n",
" rew_air_time *= (\n", " rew_air_time *= (\n",
" math.normalize(commands[:2])[1] \u003e 0.05\n", " math.normalize(commands[:2])[1] > 0.05\n",
" ) # no reward for zero command\n", " ) # no reward for zero command\n",
" return rew_air_time\n", " return rew_air_time\n",
"\n", "\n",
@@ -1259,15 +1259,15 @@
" self,\n", " self,\n",
" commands: jax.Array,\n", " commands: jax.Array,\n",
" joint_angles: jax.Array,\n", " joint_angles: jax.Array,\n",
" ) -\u003e jax.Array:\n", " ) -> jax.Array:\n",
" # Penalize motion at zero commands\n", " # Penalize motion at zero commands\n",
" return jp.sum(jp.abs(joint_angles - self._default_pose)) * (\n", " return jp.sum(jp.abs(joint_angles - self._default_pose)) * (\n",
" math.normalize(commands[:2])[1] \u003c 0.1\n", " math.normalize(commands[:2])[1] < 0.1\n",
" )\n", " )\n",
"\n", "\n",
" def _reward_foot_slip(\n", " def _reward_foot_slip(\n",
" self, pipeline_state: base.State, contact_filt: jax.Array\n", " self, pipeline_state: base.State, contact_filt: jax.Array\n",
" ) -\u003e jax.Array:\n", " ) -> jax.Array:\n",
" # get velocities at feet which are offset from lower legs\n", " # get velocities at feet which are offset from lower legs\n",
" # pytype: disable=attribute-error\n", " # pytype: disable=attribute-error\n",
" pos = pipeline_state.site_xpos[self._feet_site_id] # feet position\n", " pos = pipeline_state.site_xpos[self._feet_site_id] # feet position\n",
@@ -1280,12 +1280,12 @@
" # Penalize large feet velocity for feet that are in contact with the ground.\n", " # Penalize large feet velocity for feet that are in contact with the ground.\n",
" return jp.sum(jp.square(foot_vel[:, :2]) * contact_filt.reshape((-1, 1)))\n", " return jp.sum(jp.square(foot_vel[:, :2]) * contact_filt.reshape((-1, 1)))\n",
"\n", "\n",
" def _reward_termination(self, done: jax.Array, step: jax.Array) -\u003e jax.Array:\n", " def _reward_termination(self, done: jax.Array, step: jax.Array) -> jax.Array:\n",
" return done \u0026 (step \u003c 500)\n", " return done & (step < 500)\n",
"\n", "\n",
" def render(\n", " def render(\n",
" self, trajectory: List[base.State], camera: str | None = None\n", " self, trajectory: List[base.State], camera: str | None = None\n",
" ) -\u003e Sequence[np.ndarray]:\n", " ) -> Sequence[np.ndarray]:\n",
" camera = camera or 'track'\n", " camera = camera or 'track'\n",
" return super().render(trajectory, camera=camera)\n", " return super().render(trajectory, camera=camera)\n",
"\n", "\n",
@@ -1294,7 +1294,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "pi_yrcz-Qp3W" "id": "pi_yrcz-Qp3W"
}, },
@@ -1317,7 +1317,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "cHJCbESGA7Rk" "id": "cHJCbESGA7Rk"
}, },
@@ -1355,7 +1355,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "8Yge-CGP5JoO" "id": "8Yge-CGP5JoO"
}, },
@@ -1383,7 +1383,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "VTbpEtXnEecd" "id": "VTbpEtXnEecd"
}, },
@@ -1397,7 +1397,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "HRRN-8L-BivZ" "id": "HRRN-8L-BivZ"
}, },
@@ -1442,7 +1442,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "V7jqv08X95u4" "id": "V7jqv08X95u4"
}, },
@@ -1459,12 +1459,6 @@
"gpuType": "V100", "gpuType": "V100",
"machine_shape": "hm", "machine_shape": "hm",
"private_outputs": true, "private_outputs": true,
"provenance": [
{
"file_id": "1A58SK07tnOzix53E68D0TQ2ePCTZA61f",
"timestamp": 1707342610876
}
],
"toc_visible": true "toc_visible": true
}, },
"kernelspec": { "kernelspec": {
python/LQR.ipynb
+44 -49
View File
@@ -8,11 +8,11 @@
"source": [ "source": [
"![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n", "![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n",
"\n", "\n",
"# \u003ch1\u003e\u003ccenter\u003eLQR tutorial \u003ca href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.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", "# <h1><center>LQR tutorial <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a></center></h1>\n",
"\n", "\n",
"This notebook provides an example of an LQR controller using [**MuJoCo** physics](https://github.com/google-deepmind/mujoco#readme).\n", "This notebook provides an example of an LQR controller using [**MuJoCo** physics](https://github.com/google-deepmind/mujoco#readme).\n",
"\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" "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime > Change runtime type\".\n"
] ]
}, },
{ {
@@ -30,9 +30,9 @@
"id": "_UbO9uhtBSX5" "id": "_UbO9uhtBSX5"
}, },
"source": [ "source": [
"\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eCopyright 2022 DeepMind Technologies Limited\u003c/small\u003e\u003c/p\u003e\n", "> <p><small><small>Copyright 2022 DeepMind Technologies Limited</small></p>\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", "> <p><small><small>Licensed 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 <a href=\"http://www.apache.org/licenses/LICENSE-2.0\">http://www.apache.org/licenses/LICENSE-2.0</a>.</small></small></p>\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" "> <p><small><small>Unless 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.</small></small></p>"
] ]
}, },
{ {
@@ -46,7 +46,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "Xqo7pyX-n72M" "id": "Xqo7pyX-n72M"
}, },
@@ -57,7 +57,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"cellView": "form", "cellView": "form",
"id": "IbZxYDxzoz5R" "id": "IbZxYDxzoz5R"
@@ -99,7 +99,7 @@
"try:\n", "try:\n",
" print('Checking that the installation succeeded:')\n", " print('Checking that the installation succeeded:')\n",
" import mujoco\n", " import mujoco\n",
" mujoco.MjModel.from_xml_string('\u003cmujoco/\u003e')\n", " mujoco.MjModel.from_xml_string('<mujoco/>')\n",
"except Exception as e:\n", "except Exception as e:\n",
" raise e from RuntimeError(\n", " raise e from RuntimeError(\n",
" 'Something went wrong during installation. Check the shell output above '\n", " 'Something went wrong during installation. Check the shell output above '\n",
@@ -112,7 +112,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"cellView": "form", "cellView": "form",
"id": "T5f4w3Kq2X14" "id": "T5f4w3Kq2X14"
@@ -126,7 +126,7 @@
"\n", "\n",
"# Graphics and plotting.\n", "# Graphics and plotting.\n",
"print('Installing mediapy:')\n", "print('Installing mediapy:')\n",
"!command -v ffmpeg \u003e/dev/null || (apt update \u0026\u0026 apt install -y ffmpeg)\n", "!command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
"!pip install -q mediapy\n", "!pip install -q mediapy\n",
"import mediapy as media\n", "import mediapy as media\n",
"import matplotlib.pyplot as plt\n", "import matplotlib.pyplot as plt\n",
@@ -146,7 +146,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "RecFjafkfX4V" "id": "RecFjafkfX4V"
}, },
@@ -169,7 +169,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "yftlgN0yznRe" "id": "yftlgN0yznRe"
}, },
@@ -191,7 +191,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "F6ZhQv2l0OOu" "id": "F6ZhQv2l0OOu"
}, },
@@ -215,7 +215,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "wWzbBpCBgAzE" "id": "wWzbBpCBgAzE"
}, },
@@ -239,7 +239,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "I75-J4DowklB" "id": "I75-J4DowklB"
}, },
@@ -252,12 +252,12 @@
"mujoco.mj_resetDataKeyframe(model, data, 1)\n", "mujoco.mj_resetDataKeyframe(model, data, 1)\n",
"\n", "\n",
"frames = []\n", "frames = []\n",
"while data.time \u003c DURATION:\n", "while data.time < DURATION:\n",
" # Step the simulation.\n", " # Step the simulation.\n",
" mujoco.mj_step(model, data)\n", " mujoco.mj_step(model, data)\n",
"\n", "\n",
" # Render and save frames.\n", " # Render and save frames.\n",
" if len(frames) \u003c data.time * FRAMERATE:\n", " if len(frames) < data.time * FRAMERATE:\n",
" renderer.update_scene(data)\n", " renderer.update_scene(data)\n",
" pixels = renderer.render()\n", " pixels = renderer.render()\n",
" frames.append(pixels)\n", " frames.append(pixels)\n",
@@ -279,7 +279,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "kFeaLU7n42Iu" "id": "kFeaLU7n42Iu"
}, },
@@ -296,7 +296,7 @@
"mujoco.mj_resetDataKeyframe(model, data, 1)\n", "mujoco.mj_resetDataKeyframe(model, data, 1)\n",
"\n", "\n",
"frames = []\n", "frames = []\n",
"while data.time \u003c DURATION:\n", "while data.time < DURATION:\n",
" # Set control vector.\n", " # Set control vector.\n",
" data.ctrl = np.random.randn(model.nu)\n", " data.ctrl = np.random.randn(model.nu)\n",
"\n", "\n",
@@ -304,7 +304,7 @@
" mujoco.mj_step(model, data)\n", " mujoco.mj_step(model, data)\n",
"\n", "\n",
" # Render and save frames.\n", " # Render and save frames.\n",
" if len(frames) \u003c data.time * FRAMERATE:\n", " if len(frames) < data.time * FRAMERATE:\n",
" # Set the lookat point to the humanoid's center of mass.\n", " # Set the lookat point to the humanoid's center of mass.\n",
" camera.lookat = data.body('torso').subtree_com\n", " camera.lookat = data.body('torso').subtree_com\n",
"\n", "\n",
@@ -405,7 +405,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "8Q6Ceuf7ZHGQ" "id": "8Q6Ceuf7ZHGQ"
}, },
@@ -431,7 +431,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "2eN8b1EZ5GO9" "id": "2eN8b1EZ5GO9"
}, },
@@ -483,7 +483,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "qaw4gxg46h2G" "id": "qaw4gxg46h2G"
}, },
@@ -510,7 +510,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "a1PF_yHdPvLl" "id": "a1PF_yHdPvLl"
}, },
@@ -534,7 +534,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "dDLihz5hk9Wt" "id": "dDLihz5hk9Wt"
}, },
@@ -556,7 +556,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "8cQpEF7MmDPI" "id": "8cQpEF7MmDPI"
}, },
@@ -571,12 +571,12 @@
"data.ctrl = ctrl0\n", "data.ctrl = ctrl0\n",
"\n", "\n",
"frames = []\n", "frames = []\n",
"while data.time \u003c DURATION:\n", "while data.time < DURATION:\n",
" # Step the simulation.\n", " # Step the simulation.\n",
" mujoco.mj_step(model, data)\n", " mujoco.mj_step(model, data)\n",
"\n", "\n",
" # Render and save frames.\n", " # Render and save frames.\n",
" if len(frames) \u003c data.time * FRAMERATE:\n", " if len(frames) < data.time * FRAMERATE:\n",
" # Set the lookat point to the humanoid's center of mass.\n", " # Set the lookat point to the humanoid's center of mass.\n",
" camera.lookat = data.body('torso').subtree_com\n", " camera.lookat = data.body('torso').subtree_com\n",
" renderer.update_scene(data, camera)\n", " renderer.update_scene(data, camera)\n",
@@ -608,7 +608,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "SxOeHSf1nspy" "id": "SxOeHSf1nspy"
}, },
@@ -631,7 +631,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "9LEK4MCHkH15" "id": "9LEK4MCHkH15"
}, },
@@ -671,7 +671,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "S731eD3Jm1mJ" "id": "S731eD3Jm1mJ"
}, },
@@ -706,7 +706,7 @@
"id": "OeHYWdQdm-vE" "id": "OeHYWdQdm-vE"
}, },
"source": [ "source": [
"We are now ready to construct the Q matrix. Note that the coefficient of the balancing term is quite high. This is due to 3 seperate reasons:\n", "We are now ready to construct the Q matrix. Note that the coefficient of the balancing term is quite high. This is due to 3 separate reasons:\n",
"- It's the thing we care about most. Balancing means keeping the CoM over the foot.\n", "- It's the thing we care about most. Balancing means keeping the CoM over the foot.\n",
"- We have less control authority over the CoM (relative to body joints).\n", "- We have less control authority over the CoM (relative to body joints).\n",
"- In the balancing context, units of length are \"bigger\". If the knee bends by 0.1 radians (≈6°), we can probably still recover. If the CoM position is 10cm sideways from the foot position, we are likely on our way to the floor." "- In the balancing context, units of length are \"bigger\". If the knee bends by 0.1 radians (≈6°), we can probably still recover. If the CoM position is 10cm sideways from the foot position, we are likely on our way to the floor."
@@ -714,7 +714,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "reIA8___o3Z4" "id": "reIA8___o3Z4"
}, },
@@ -752,7 +752,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "NB4ZStYrpx1B" "id": "NB4ZStYrpx1B"
}, },
@@ -782,7 +782,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "azjsWl4v_K16" "id": "azjsWl4v_K16"
}, },
@@ -810,7 +810,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "Z_57VMUDpGrj" "id": "Z_57VMUDpGrj"
}, },
@@ -828,7 +828,7 @@
"dq = np.zeros(model.nv)\n", "dq = np.zeros(model.nv)\n",
"\n", "\n",
"frames = []\n", "frames = []\n",
"while data.time \u003c DURATION:\n", "while data.time < DURATION:\n",
" # Get state difference dx.\n", " # Get state difference dx.\n",
" mujoco.mj_differentiatePos(model, dq, 1, qpos0, data.qpos)\n", " mujoco.mj_differentiatePos(model, dq, 1, qpos0, data.qpos)\n",
" dx = np.hstack((dq, data.qvel)).T\n", " dx = np.hstack((dq, data.qvel)).T\n",
@@ -840,7 +840,7 @@
" mujoco.mj_step(model, data)\n", " mujoco.mj_step(model, data)\n",
"\n", "\n",
" # Render and save frames.\n", " # Render and save frames.\n",
" if len(frames) \u003c data.time * FRAMERATE:\n", " if len(frames) < data.time * FRAMERATE:\n",
" renderer.update_scene(data)\n", " renderer.update_scene(data)\n",
" pixels = renderer.render()\n", " pixels = renderer.render()\n",
" frames.append(pixels)\n", " frames.append(pixels)\n",
@@ -865,7 +865,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": null, "execution_count": 0,
"metadata": { "metadata": {
"id": "zJmbOJMurRna" "id": "zJmbOJMurRna"
}, },
@@ -891,9 +891,9 @@
"model.vis.map.force = 0.01\n", "model.vis.map.force = 0.01\n",
"\n", "\n",
"# Define smooth orbiting function.\n", "# Define smooth orbiting function.\n",
"def unit_smooth(normalised_time: float) -\u003e float:\n", "def unit_smooth(normalised_time: float) -> float:\n",
" return 1 - np.cos(normalised_time*2*np.pi)\n", " return 1 - np.cos(normalised_time*2*np.pi)\n",
"def azimuth(time: float) -\u003e float:\n", "def azimuth(time: float) -> float:\n",
" return 100 + unit_smooth(data.time/DURATION) * TOTAL_ROTATION\n", " return 100 + unit_smooth(data.time/DURATION) * TOTAL_ROTATION\n",
"\n", "\n",
"# Precompute some noise.\n", "# Precompute some noise.\n",
@@ -917,7 +917,7 @@
"\n", "\n",
"frames = []\n", "frames = []\n",
"step = 0\n", "step = 0\n",
"while data.time \u003c DURATION:\n", "while data.time < DURATION:\n",
" # Get state difference dx.\n", " # Get state difference dx.\n",
" mujoco.mj_differentiatePos(model, dq, 1, qpos0, data.qpos)\n", " mujoco.mj_differentiatePos(model, dq, 1, qpos0, data.qpos)\n",
" dx = np.hstack((dq, data.qvel)).T\n", " dx = np.hstack((dq, data.qvel)).T\n",
@@ -933,7 +933,7 @@
" mujoco.mj_step(model, data)\n", " mujoco.mj_step(model, data)\n",
"\n", "\n",
" # Render and save frames.\n", " # Render and save frames.\n",
" if len(frames) \u003c data.time * FRAMERATE:\n", " if len(frames) < data.time * FRAMERATE:\n",
" camera.azimuth = azimuth(data.time)\n", " camera.azimuth = azimuth(data.time)\n",
" renderer.update_scene(data, camera, scene_option)\n", " renderer.update_scene(data, camera, scene_option)\n",
" pixels = renderer.render()\n", " pixels = renderer.render()\n",
@@ -949,12 +949,7 @@
"collapsed_sections": [ "collapsed_sections": [
"LBAvTJ0xHKy7" "LBAvTJ0xHKy7"
], ],
"last_runtime": { "private_outputs": true
"build_target": "",
"kind": "local"
},
"private_outputs": true,
"provenance": []
}, },
"kernelspec": { "kernelspec": {
"display_name": "Python 3", "display_name": "Python 3",
python/tutorial.ipynb
+341 -342
View File
File diff suppressed because it is too large Load Diff