Mujoco-Warp notebook examples

examples/machine_learning/.gitignore

@@ -0,0 +1,2 @@
+*.msgpack
+*.pt

examples/machine_learning/stretch4_mujoco_warp_policy_example.ipynb

@@ -0,0 +1,234 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Policy Training with MuJoCo Warp\n",
+    "\n",
+    "This notebook provides a skeleton for training a reinforcement learning policy using the Stretch 4 Mujoco model, Robocasa and MuJoCo Warp in Python. \n",
+    "MuJoCo Warp (MJWarp) enables high-throughput, parallelized physics simulations directly on NVIDIA GPUs, drastically reducing RL model training time."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Environment Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-05-08T01:16:01.973014Z",
+     "iopub.status.busy": "2026-05-08T01:16:01.972955Z",
+     "iopub.status.idle": "2026-05-08T01:16:02.881691Z",
+     "shell.execute_reply": "2026-05-08T01:16:02.880973Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "First, follow the README at the root of the stretch4_mujoco repo \n",
+    "to install the dependencies for Stretch4 Mujoco and Robocasa.\n",
+    "Afterwards, come back to run this cell.\n",
+    "\"\"\"\n",
+    "!uv pip install mujoco-warp\n",
+    "!uv pip install -U \"jax[cuda12]\" optax flax\n",
+    "!uv pip install numpy==2.2.5\n",
+    "!uv pip install mujoco==3.3.1\n",
+    "!uv pip install mediapy"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-05-08T01:16:02.883778Z",
+     "iopub.status.busy": "2026-05-08T01:16:02.883612Z",
+     "iopub.status.idle": "2026-05-08T01:16:03.371724Z",
+     "shell.execute_reply": "2026-05-08T01:16:03.367579Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from stretch4_mujoco_warp_policy_example import (\n",
+    "    setup_environment,\n",
+    "    get_default_config,\n",
+    "    generate_robocasa_model,\n",
+    "    load_mujoco_warp_model,\n",
+    "    PPOTrainer,\n",
+    "    run_training,\n",
+    "    save_params,\n",
+    "    run_viewer,\n",
+    "    render_video\n",
+    ")\n",
+    "\n",
+    "setup_environment()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Load the Stretch4 Mujoco Environment in Robocasa"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = get_default_config()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mj_model, target_object_body_name = generate_robocasa_model()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Object pre-grasp policy example\n",
+    "\n",
+    "The following cell shows an example of a PPO policy that uses the `obj_main` object in mujoco as the target object for Stretch 4 to attempt to pre-grasp (position the gripper near the object)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load the Mujoco model into Mujoco-Warp"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "warp_model, warp_data = load_mujoco_warp_model(mj_model, config)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## PPO Policy and Loss Functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-05-08T01:16:03.375183Z",
+     "iopub.status.busy": "2026-05-08T01:16:03.374902Z",
+     "iopub.status.idle": "2026-05-08T01:17:38.054443Z",
+     "shell.execute_reply": "2026-05-08T01:17:38.053818Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "trainer = PPOTrainer(mj_model, warp_model, warp_data, target_object_body_name, config)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Run Training\n",
+    "\n",
+    "This takes about 30 minutes"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "params, training_frames = run_training(trainer)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Saving Training data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "filename = save_params(params)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Loading saved training data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "run_viewer(mj_model, target_object_body_name, params)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Rendering videos"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "render_video(mj_model, target_object_body_name, params)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}