{ "cells": [ { "cell_type": "markdown", "id": "318a00e9", "metadata": {}, "source": [ "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tayalmanan28/MuJoCo-Tutorial/blob/main/tutorial/03_simulation.ipynb)\n", "\n", "# Tutorial 3: Simulation\n", "\n", "Understand `mj_step`, `mj_forward`, timesteps, and the simulation pipeline." ] }, { "cell_type": "code", "execution_count": null, "id": "363d27cc", "metadata": {}, "outputs": [], "source": [ "!pip install -q mujoco mediapy matplotlib" ] }, { "cell_type": "code", "execution_count": null, "id": "aaecfaba", "metadata": {}, "outputs": [], "source": [ "import mujoco\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import mediapy as media" ] }, { "cell_type": "markdown", "id": "6127660c", "metadata": {}, "source": [ "## The Simulation Loop\n", "\n", "MuJoCo's main function is `mj_step(model, data)` which advances the state by one timestep:\n", "\n", "$$x_{t+h} = f(x_t, u_t)$$\n", "\n", "where $x = (q, \\dot{q})$ is the state and $u$ is the control." ] }, { "cell_type": "code", "execution_count": null, "id": "d5410dc6", "metadata": {}, "outputs": [], "source": [ "xml = \"\"\"\n", "\n", " \n", "\"\"\"\n", "\n", "model = mujoco.MjModel.from_xml_string(xml)\n", "data = mujoco.MjData(model)\n", "\n", "print(f\"Timestep: {model.opt.timestep}s\")\n", "print(f\"Initial height: {data.qpos[2]:.3f}\")\n", "\n", "# Simulate for 1 second\n", "times, heights = [], []\n", "while data.time < 1.0:\n", " mujoco.mj_step(model, data)\n", " times.append(data.time)\n", " heights.append(data.qpos[2])\n", "\n", "plt.figure(figsize=(10, 4))\n", "plt.plot(times, heights)\n", "plt.xlabel('Time (s)'); plt.ylabel('Height (m)')\n", "plt.title('Ball dropping and bouncing')\n", "plt.grid(True); plt.show()" ] }, { "cell_type": "markdown", "id": "51f03701", "metadata": {}, "source": [ "## `mj_step` vs `mj_forward`\n", "\n", "- **`mj_step`**: Advances time. Computes forces → integrates → updates state.\n", "- **`mj_forward`**: Computes all derived quantities **without** advancing time.\n", "\n", "Use `mj_forward` when you need positions/velocities after setting `qpos`/`qvel` manually." ] }, { "cell_type": "code", "execution_count": null, "id": "2bdada12", "metadata": {}, "outputs": [], "source": [ "# Reset and manually set position\n", "mujoco.mj_resetData(model, data)\n", "data.qpos[2] = 5.0 # set height to 5m\n", "\n", "# mj_forward propagates derived quantities\n", "mujoco.mj_forward(model, data)\n", "print(f\"Ball position: {data.qpos[:3]}\")\n", "print(f\"Ball velocity: {data.qvel[:3]}\")\n", "print(f\"Time did NOT advance: {data.time}\")" ] }, { "cell_type": "markdown", "id": "29ed0624", "metadata": {}, "source": [ "## Timestep and Stability\n", "\n", "Smaller timesteps = more accurate but slower. MuJoCo's default is 2ms.\n", "Let's compare:" ] }, { "cell_type": "code", "execution_count": null, "id": "a9ce2040", "metadata": {}, "outputs": [], "source": [ "plt.figure(figsize=(10, 4))\n", "\n", "for dt in [0.001, 0.005, 0.01]:\n", " model_dt = mujoco.MjModel.from_xml_string(xml)\n", " model_dt.opt.timestep = dt\n", " data_dt = mujoco.MjData(model_dt)\n", " \n", " times, heights = [], []\n", " while data_dt.time < 1.0:\n", " mujoco.mj_step(model_dt, data_dt)\n", " times.append(data_dt.time)\n", " heights.append(data_dt.qpos[2])\n", " \n", " plt.plot(times, heights, label=f'dt={dt}s ({len(times)} steps)')\n", "\n", "plt.xlabel('Time (s)'); plt.ylabel('Height (m)')\n", "plt.title('Effect of timestep on simulation')\n", "plt.legend(); plt.grid(True); plt.show()" ] }, { "cell_type": "markdown", "id": "242bd9e5", "metadata": {}, "source": [ "## Energy & Momentum Tracking" ] }, { "cell_type": "code", "execution_count": null, "id": "800ab5fb", "metadata": {}, "outputs": [], "source": [ "model_e = mujoco.MjModel.from_xml_string(xml)\n", "model_e.opt.timestep = 0.001\n", "data_e = mujoco.MjData(model_e)\n", "\n", "# Disable contacts to see energy conservation\n", "model_e.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_CONTACT\n", "\n", "times, kinetic, potential = [], [], []\n", "while data_e.time < 0.45: # before hitting ground\n", " mujoco.mj_step(model_e, data_e)\n", " times.append(data_e.time)\n", " # KE = 0.5 * m * v^2\n", " v = np.linalg.norm(data_e.qvel[:3])\n", " kinetic.append(0.5 * 1.0 * v**2)\n", " # PE = m * g * h\n", " potential.append(1.0 * 9.81 * data_e.qpos[2])\n", "\n", "total = np.array(kinetic) + np.array(potential)\n", "\n", "plt.figure(figsize=(10, 4))\n", "plt.plot(times, kinetic, label='Kinetic')\n", "plt.plot(times, potential, label='Potential')\n", "plt.plot(times, total, 'k--', label='Total')\n", "plt.xlabel('Time (s)'); plt.ylabel('Energy (J)')\n", "plt.title('Energy conservation (no contacts)')\n", "plt.legend(); plt.grid(True); plt.show()\n", "print(f\"Energy drift: {(total[-1] - total[0])/total[0]*100:.6f}%\")" ] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }