Upload opentrack.ipynb

samples/opentrack.ipynb

@@ -0,0 +1,493 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# OpenTrack - Humanoid Motion Tracking Demo\n",
+    "\n",
+    "This notebook demonstrates OpenTrack, an open-source humanoid motion tracking system using MuJoCo.\n",
+    "\n",
+    "**Note**: Training can be resource-intensive. We'll use debug mode for quick testing."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Setup Environment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import subprocess\n",
+    "import time\n",
+    "import glob\n",
+    "from pathlib import Path\n",
+    "from IPython.display import Video, display, HTML\n",
+    "import threading\n",
+    "import queue\n",
+    "\n",
+    "print(\"Environment setup complete!\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Clone OpenTrack Repository"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Clone the repository if not already cloned\n",
+    "if not os.path.exists('OpenTrack'):\n",
+    "    !git clone https://github.com/GalaxyGeneralRobotics/OpenTrack.git\n",
+    "    print(\"✓ Repository cloned successfully\")\n",
+    "else:\n",
+    "    print(\"✓ Repository already exists\")\n",
+    "\n",
+    "os.chdir('OpenTrack')\n",
+    "print(f\"Current directory: {os.getcwd()}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Install Dependencies"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install PyTorch (CPU version for compatibility)\n",
+    "!pip install -q torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cpu\n",
+    "\n",
+    "# Install OpenTrack requirements\n",
+    "!pip install -q -r requirements.txt\n",
+    "\n",
+    "# Install additional packages for video handling\n",
+    "!pip install -q imageio imageio-ffmpeg\n",
+    "\n",
+    "print(\"✓ All dependencies installed\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Download Motion Capture Data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import snapshot_download\n",
+    "\n",
+    "# Create directory structure\n",
+    "mocap_dir = Path(\"data/mocap/lafan1/UnitreeG1\")\n",
+    "mocap_dir.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "repo_id = \"robfiras/loco-mujoco-datasets\"\n",
+    "\n",
+    "print(\"Downloading all mocap data from Lafan1/mocap/UnitreeG1...\")\n",
+    "print(\"This will download all .npz files concurrently.\\n\")\n",
+    "\n",
+    "try:\n",
+    "    # Use snapshot_download with allow_patterns to download only the files we need\n",
+    "    # This is much more efficient than downloading files one by one\n",
+    "    snapshot_path = snapshot_download(\n",
+    "        repo_id=repo_id,\n",
+    "        repo_type=\"dataset\",\n",
+    "        allow_patterns=\"Lafan1/mocap/UnitreeG1/*.npz\",\n",
+    "        local_dir=\".\",\n",
+    "        local_dir_use_symlinks=False\n",
+    "    )\n",
+    "    \n",
+    "    print(f\"\\n✓ Download complete! Files saved to: {snapshot_path}\")\n",
+    "    \n",
+    "    # Verify files\n",
+    "    npz_files = list(mocap_dir.glob(\"*.npz\"))\n",
+    "    print(f\"✓ Found {len(npz_files)} .npz files in {mocap_dir}\")\n",
+    "    \n",
+    "    if npz_files:\n",
+    "        print(\"\\nSample files:\")\n",
+    "        for f in sorted(npz_files)[:10]:  # Show first 10 files\n",
+    "            print(f\"  - {f.name}\")\n",
+    "        if len(npz_files) > 10:\n",
+    "            print(f\"  ... and {len(npz_files) - 10} more files\")\n",
+    "    \n",
+    "except Exception as e:\n",
+    "    print(f\"⚠ Error downloading mocap data: {e}\")\n",
+    "    print(\"\\nYou may need to download manually from:\")\n",
+    "    print(\"https://huggingface.co/datasets/robfiras/loco-mujoco-datasets/tree/main/Lafan1/mocap/UnitreeG1\")\n",
+    "    print(\"\\nOr check if you need to authenticate:\")\n",
+    "    print(\"  huggingface-cli login\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Helper Functions for Background Process Management"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def run_command_with_output(cmd, description=\"Running command\"):\n",
+    "    \"\"\"\n",
+    "    Run a command and display output in real-time\n",
+    "    \"\"\"\n",
+    "    print(f\"\\n{'='*60}\")\n",
+    "    print(f\"{description}\")\n",
+    "    print(f\"Command: {' '.join(cmd)}\")\n",
+    "    print(f\"{'='*60}\\n\")\n",
+    "    \n",
+    "    process = subprocess.Popen(\n",
+    "        cmd,\n",
+    "        stdout=subprocess.PIPE,\n",
+    "        stderr=subprocess.STDOUT,\n",
+    "        text=True,\n",
+    "        bufsize=1,\n",
+    "        universal_newlines=True\n",
+    "    )\n",
+    "    \n",
+    "    # Read output line by line\n",
+    "    for line in process.stdout:\n",
+    "        print(line, end='')\n",
+    "    \n",
+    "    process.wait()\n",
+    "    \n",
+    "    if process.returncode == 0:\n",
+    "        print(f\"\\n✓ {description} completed successfully\")\n",
+    "    else:\n",
+    "        print(f\"\\n⚠ {description} exited with code {process.returncode}\")\n",
+    "    \n",
+    "    return process.returncode\n",
+    "\n",
+    "\n",
+    "def find_latest_experiment(exp_name_pattern):\n",
+    "    \"\"\"\n",
+    "    Find the latest experiment folder matching the pattern\n",
+    "    \"\"\"\n",
+    "    logs_dir = Path(\"logs\")\n",
+    "    if not logs_dir.exists():\n",
+    "        return None\n",
+    "    \n",
+    "    # Find all matching experiments\n",
+    "    experiments = sorted(\n",
+    "        [d for d in logs_dir.iterdir() if d.is_dir() and exp_name_pattern in d.name],\n",
+    "        key=lambda x: x.stat().st_mtime,\n",
+    "        reverse=True\n",
+    "    )\n",
+    "    \n",
+    "    return experiments[0].name if experiments else None\n",
+    "\n",
+    "\n",
+    "def find_generated_videos(output_dir=\"videos\"):\n",
+    "    \"\"\"\n",
+    "    Find all generated video files\n",
+    "    \"\"\"\n",
+    "    video_dir = Path(output_dir)\n",
+    "    if not video_dir.exists():\n",
+    "        # Try alternative locations\n",
+    "        alternative_dirs = [\".\", \"logs\", \"outputs\"]\n",
+    "        for alt_dir in alternative_dirs:\n",
+    "            alt_path = Path(alt_dir)\n",
+    "            videos = list(alt_path.glob(\"**/*.mp4\")) + list(alt_path.glob(\"**/*.gif\"))\n",
+    "            if videos:\n",
+    "                return sorted(videos, key=lambda x: x.stat().st_mtime, reverse=True)\n",
+    "        return []\n",
+    "    \n",
+    "    videos = list(video_dir.glob(\"*.mp4\")) + list(video_dir.glob(\"*.gif\"))\n",
+    "    return sorted(videos, key=lambda x: x.stat().st_mtime, reverse=True)\n",
+    "\n",
+    "print(\"✓ Helper functions loaded\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Quick Training (Debug Mode)\n",
+    "\n",
+    "We'll run a quick training session in debug mode. This won't produce a well-trained model but will verify everything works."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Run quick training\n",
+    "cmd = [\n",
+    "    'python', 'train_policy.py',\n",
+    "    '--exp_name', 'debug',\n",
+    "    '--terrain_type', 'flat_terrain'\n",
+    "]\n",
+    "\n",
+    "return_code = run_command_with_output(cmd, \"Training OpenTrack (debug mode)\")\n",
+    "\n",
+    "# Find the experiment folder\n",
+    "exp_folder = find_latest_experiment('debug')\n",
+    "if exp_folder:\n",
+    "    print(f\"\\n✓ Training completed! Experiment: {exp_folder}\")\n",
+    "else:\n",
+    "    print(\"\\n⚠ Could not find experiment folder\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 7. Convert Checkpoint (Brax → PyTorch)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Get the latest experiment name\n",
+    "exp_folder = find_latest_experiment('debug')\n",
+    "\n",
+    "if exp_folder:\n",
+    "    print(f\"Converting checkpoint for: {exp_folder}\")\n",
+    "    \n",
+    "    cmd = [\n",
+    "        'python', 'brax2torch.py',\n",
+    "        '--exp_name', exp_folder\n",
+    "    ]\n",
+    "    \n",
+    "    return_code = run_command_with_output(cmd, \"Converting Brax checkpoint to PyTorch\")\n",
+    "else:\n",
+    "    print(\"⚠ No experiment found. Run training first.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. Generate Videos (Headless Rendering)\n",
+    "\n",
+    "This will run the policy and generate videos using MuJoCo's headless renderer."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Get the latest experiment name\n",
+    "exp_folder = find_latest_experiment('debug')\n",
+    "\n",
+    "if exp_folder:\n",
+    "    print(f\"Generating videos for: {exp_folder}\")\n",
+    "    \n",
+    "    # Use --use_renderer for headless video generation (NOT --use_viewer)\n",
+    "    cmd = [\n",
+    "        'python', 'play_policy.py',\n",
+    "        '--exp_name', exp_folder,\n",
+    "        '--use_renderer',\n",
+    "        # '--play_ref_motion',  # Uncomment to also show reference motion\n",
+    "    ]\n",
+    "    \n",
+    "    return_code = run_command_with_output(cmd, \"Generating videos with MuJoCo renderer\")\n",
+    "    \n",
+    "    print(\"\\n\" + \"=\"*60)\n",
+    "    print(\"Searching for generated videos...\")\n",
+    "    print(\"=\"*60)\n",
+    "    \n",
+    "    # Wait a moment for files to be written\n",
+    "    time.sleep(2)\n",
+    "    \n",
+    "    # Find generated videos\n",
+    "    videos = find_generated_videos()\n",
+    "    \n",
+    "    if videos:\n",
+    "        print(f\"\\n✓ Found {len(videos)} video(s):\")\n",
+    "        for video in videos:\n",
+    "            print(f\"  - {video}\")\n",
+    "    else:\n",
+    "        print(\"\\n⚠ No videos found. Checking alternative locations...\")\n",
+    "        # Search more broadly\n",
+    "        all_videos = list(Path(\".\").rglob(\"*.mp4\")) + list(Path(\".\").rglob(\"*.gif\"))\n",
+    "        if all_videos:\n",
+    "            print(f\"Found {len(all_videos)} video(s) in project:\")\n",
+    "            for video in all_videos[:10]:  # Show first 10\n",
+    "                print(f\"  - {video}\")\n",
+    "else:\n",
+    "    print(\"⚠ No experiment found. Run training first.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 9. Display Generated Videos"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Find and display all generated videos\n",
+    "videos = find_generated_videos()\n",
+    "\n",
+    "if not videos:\n",
+    "    # Try alternative search\n",
+    "    videos = list(Path(\".\").rglob(\"*.mp4\")) + list(Path(\".\").rglob(\"*.gif\"))\n",
+    "\n",
+    "if videos:\n",
+    "    print(f\"Displaying {len(videos)} video(s):\\n\")\n",
+    "    \n",
+    "    for i, video_path in enumerate(videos[:5]):  # Display first 5 videos\n",
+    "        print(f\"\\n{'='*60}\")\n",
+    "        print(f\"Video {i+1}: {video_path.name}\")\n",
+    "        print(f\"{'='*60}\")\n",
+    "        \n",
+    "        try:\n",
+    "            if video_path.suffix == '.mp4':\n",
+    "                display(Video(str(video_path), width=800, embed=True))\n",
+    "            elif video_path.suffix == '.gif':\n",
+    "                display(HTML(f'<img src=\"{video_path}\" width=\"800\">'))\n",
+    "        except Exception as e:\n",
+    "            print(f\"⚠ Error displaying video: {e}\")\n",
+    "            print(f\"You can download it from: {video_path}\")\n",
+    "else:\n",
+    "    print(\"⚠ No videos found.\")\n",
+    "    print(\"\\nPossible reasons:\")\n",
+    "    print(\"1. Training didn't complete successfully\")\n",
+    "    print(\"2. Checkpoint conversion failed\")\n",
+    "    print(\"3. Video generation failed\")\n",
+    "    print(\"\\nCheck the output of previous cells for errors.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 10. Optional: Generate Rough Terrain\n",
+    "\n",
+    "If you want to test on rough terrain, run this cell first to generate terrain data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Generate rough terrain using Perlin noise\n",
+    "cmd = ['python', 'generate_terrain.py']\n",
+    "return_code = run_command_with_output(cmd, \"Generating rough terrain\")\n",
+    "\n",
+    "print(\"\\n✓ Terrain generated! You can now train with --terrain_type rough_terrain\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 11. Optional: Full Training (Takes Much Longer)\n",
+    "\n",
+    "⚠️ Warning: This will take significant time and resources. Only run if you have GPU access and time."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Full training (uncomment to run)\n",
+    "# cmd = [\n",
+    "#     'python', 'train_policy.py',\n",
+    "#     '--exp_name', 'flat_terrain_full',\n",
+    "#     '--terrain_type', 'flat_terrain'\n",
+    "# ]\n",
+    "# \n",
+    "# return_code = run_command_with_output(cmd, \"Full training on flat terrain\")\n",
+    "# \n",
+    "# # Then convert and play\n",
+    "# exp_folder = find_latest_experiment('flat_terrain_full')\n",
+    "# if exp_folder:\n",
+    "#     !python brax2torch.py --exp_name {exp_folder}\n",
+    "#     !python play_policy.py --exp_name {exp_folder} --use_renderer\n",
+    "\n",
+    "print(\"Full training cell ready (currently commented out)\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Summary\n",
+    "\n",
+    "This notebook demonstrates:\n",
+    "1. ✅ Setting up OpenTrack in a Jupyter environment\n",
+    "2. ✅ Running training (debug mode for quick testing)\n",
+    "3. ✅ Converting Brax checkpoints to PyTorch\n",
+    "4. ✅ Generating videos using headless MuJoCo renderer\n",
+    "5. ✅ Displaying videos in the notebook\n",
+    "\n",
+    "**Next Steps:**\n",
+    "- Download more mocap files for better training\n",
+    "- Run full training with GPU support\n",
+    "- Test on rough terrain\n",
+    "- Experiment with reference motion playback\n",
+    "\n",
+    "**Troubleshooting:**\n",
+    "- If videos aren't generated, check that `--use_renderer` flag is used (not `--use_viewer`)\n",
+    "- Ensure MuJoCo can run headless (may need `xvfb` on some systems)\n",
+    "- Check `logs/` directory for experiment outputs"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}