{ "cells": [ { "cell_type": "markdown", "id": "38068673", "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/08_gymnasium.ipynb)\n", "\n", "# Tutorial 8: MuJoCo + Gymnasium (RL Environments)\n", "\n", "Use MuJoCo environments for reinforcement learning with Gymnasium." ] }, { "cell_type": "code", "execution_count": null, "id": "ef44a572", "metadata": {}, "outputs": [], "source": [ "!pip install -q mujoco gymnasium[mujoco] mediapy matplotlib" ] }, { "cell_type": "code", "execution_count": null, "id": "eba07007", "metadata": {}, "outputs": [], "source": [ "import gymnasium as gym\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import mediapy as media" ] }, { "cell_type": "markdown", "id": "7c1279c3", "metadata": {}, "source": [ "## Built-in MuJoCo Environments\n", "\n", "Gymnasium ships with classic MuJoCo control tasks:" ] }, { "cell_type": "code", "execution_count": null, "id": "abc4cd85", "metadata": {}, "outputs": [], "source": [ "# List available MuJoCo environments\n", "mujoco_envs = [e for e in gym.envs.registry.keys() if 'v5' in e and any(\n", " name in e for name in ['Ant', 'HalfCheetah', 'Hopper', 'Humanoid', 'Walker', 'Swimmer', 'Reacher', 'Pendulum', 'Inverted']\n", ")]\n", "print(\"MuJoCo v5 environments:\")\n", "for e in sorted(mujoco_envs)[:12]:\n", " print(f\" {e}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "07d18a13", "metadata": {}, "outputs": [], "source": [ "# Create environment with RGB rendering\n", "env = gym.make('InvertedPendulum-v5', render_mode='rgb_array')\n", "obs, info = env.reset(seed=42)\n", "\n", "print(f\"Observation shape: {obs.shape}\")\n", "print(f\"Action space: {env.action_space}\")\n", "print(f\"Observation: {obs}\")" ] }, { "cell_type": "markdown", "id": "dce4e0d0", "metadata": {}, "source": [ "## Running an Episode" ] }, { "cell_type": "code", "execution_count": null, "id": "d63c4a6a", "metadata": {}, "outputs": [], "source": [ "frames = []\n", "obs, info = env.reset(seed=42)\n", "total_reward = 0\n", "\n", "for step in range(200):\n", " # Simple proportional controller (not RL — just to show the env working)\n", " # obs[1] is pole angle, obs[3] is pole angular velocity\n", " action = np.array([10.0 * obs[1] + 2.0 * obs[3]])\n", " action = np.clip(action, env.action_space.low, env.action_space.high)\n", " \n", " obs, reward, terminated, truncated, info = env.step(action)\n", " total_reward += reward\n", " \n", " if step % 3 == 0:\n", " frames.append(env.render())\n", " \n", " if terminated or truncated:\n", " break\n", "\n", "env.close()\n", "print(f\"Episode: {step+1} steps, reward: {total_reward:.1f}\")\n", "media.show_video(frames, fps=20)" ] }, { "cell_type": "markdown", "id": "f2e7665b", "metadata": {}, "source": [ "## Comparing Environments" ] }, { "cell_type": "code", "execution_count": null, "id": "c3664505", "metadata": {}, "outputs": [], "source": [ "env_names = ['HalfCheetah-v5', 'Ant-v5', 'Hopper-v5', 'Walker2d-v5']\n", "images = []\n", "\n", "for name in env_names:\n", " env = gym.make(name, render_mode='rgb_array')\n", " env.reset(seed=0)\n", " # Step a few times so it's not in default pose\n", " for _ in range(50):\n", " env.step(env.action_space.sample())\n", " images.append(env.render())\n", " env.close()\n", "\n", "media.show_images(images, titles=env_names)" ] }, { "cell_type": "markdown", "id": "7d627d38", "metadata": {}, "source": [ "## Creating a Custom MuJoCo Gymnasium Environment\n", "\n", "Subclass `gymnasium.Env` and use MuJoCo internally:" ] }, { "cell_type": "code", "execution_count": null, "id": "b0279096", "metadata": {}, "outputs": [], "source": [ "import mujoco\n", "from gymnasium import spaces\n", "\n", "class SimplePendulumEnv(gym.Env):\n", " \"\"\"Custom MuJoCo pendulum environment.\"\"\"\n", " metadata = {'render_modes': ['rgb_array'], 'render_fps': 30}\n", " \n", " XML = \"\"\"\n", " \n", " \n", " \"\"\"\n", " \n", " def __init__(self, render_mode=None):\n", " self.model = mujoco.MjModel.from_xml_string(self.XML)\n", " self.data = mujoco.MjData(self.model)\n", " self.render_mode = render_mode\n", " \n", " # Observation: [cos(angle), sin(angle), angular_velocity]\n", " self.observation_space = spaces.Box(-np.inf, np.inf, shape=(3,), dtype=np.float32)\n", " self.action_space = spaces.Box(-3.0, 3.0, shape=(1,), dtype=np.float32)\n", " \n", " if render_mode == 'rgb_array':\n", " self.renderer = mujoco.Renderer(self.model, height=240, width=320)\n", " \n", " def reset(self, seed=None, **kwargs):\n", " super().reset(seed=seed)\n", " mujoco.mj_resetData(self.model, self.data)\n", " # Random initial angle\n", " self.data.qpos[0] = self.np_random.uniform(-np.pi, np.pi)\n", " mujoco.mj_forward(self.model, self.data)\n", " return self._get_obs(), {}\n", " \n", " def step(self, action):\n", " self.data.ctrl[0] = action[0]\n", " mujoco.mj_step(self.model, self.data)\n", " \n", " obs = self._get_obs()\n", " # Reward: upright position\n", " angle = self.data.qpos[0]\n", " reward = np.cos(angle) - 0.01 * action[0]**2\n", " \n", " return obs, float(reward), False, False, {}\n", " \n", " def _get_obs(self):\n", " angle = self.data.qpos[0]\n", " return np.array([np.cos(angle), np.sin(angle), self.data.qvel[0]], dtype=np.float32)\n", " \n", " def render(self):\n", " if self.render_mode == 'rgb_array':\n", " self.renderer.update_scene(self.data)\n", " return self.renderer.render()\n", "\n", "# Test it\n", "env = SimplePendulumEnv(render_mode='rgb_array')\n", "obs, _ = env.reset(seed=0)\n", "print(f\"Initial obs: {obs}\")\n", "\n", "frames = []\n", "for _ in range(100):\n", " action = env.action_space.sample()\n", " obs, reward, _, _, _ = env.step(action)\n", " frames.append(env.render())\n", "\n", "media.show_video(frames, fps=30)\n", "print(\"Custom environment works!\")" ] }, { "cell_type": "markdown", "id": "8e5f1120", "metadata": {}, "source": [ "## Next Steps\n", "\n", "- Train an RL agent with [Stable-Baselines3](https://stable-baselines3.readthedocs.io/)\n", "- Use [MJX](./07_mjx.ipynb) for GPU-parallel training\n", "- Build complex environments with multiple bodies/sensors\n", "- Add reward shaping and curriculum learning" ] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }