samples/opentrack/opentrack.ipynb

{
 "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",
    "# Define project directory structure\n",
    "PROJECT_BASE = Path(\"/data/workspaces/opentrack\")\n",
    "DATASETS_DIR = PROJECT_BASE / \"datasets\"\n",
    "MODELS_DIR = PROJECT_BASE / \"models\"\n",
    "VIDEOS_DIR = PROJECT_BASE / \"videos\"\n",
    "REPO_DIR = PROJECT_BASE / \"OpenTrack\"\n",
    "\n",
    "# Create directories\n",
    "for dir_path in [PROJECT_BASE, DATASETS_DIR, MODELS_DIR, VIDEOS_DIR]:\n",
    "    dir_path.mkdir(parents=True, exist_ok=True)\n",
    "    print(f\"✓ {dir_path}\")\n",
    "\n",
    "print(\"\\n✓ 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 REPO_DIR.exists():\n",
    "    print(f\"Cloning OpenTrack to {REPO_DIR}...\")\n",
    "    subprocess.run(\n",
    "        ['git', 'clone', 'https://github.com/GalaxyGeneralRobotics/OpenTrack.git', str(REPO_DIR)],\n",
    "        check=True\n",
    "    )\n",
    "    print(\"✓ Repository cloned successfully\")\n",
    "else:\n",
    "    print(f\"✓ Repository already exists at {REPO_DIR}\")\n",
    "\n",
    "# Change to repository directory\n",
    "os.chdir(REPO_DIR)\n",
    "print(f\"\\n✓ Working 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",
    "# Define mocap directory in our datasets folder\n",
    "mocap_dir = DATASETS_DIR / \"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(f\"Target directory: {mocap_dir}\")\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",
    "    snapshot_path = snapshot_download(\n",
    "        repo_id=repo_id,\n",
    "        repo_type=\"dataset\",\n",
    "        allow_patterns=\"Lafan1/mocap/UnitreeG1/*.npz\",\n",
    "        local_dir=str(DATASETS_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",
    "    # Create symlink from OpenTrack's expected data directory to our datasets\n",
    "    opentrack_data_dir = REPO_DIR / \"data\" / \"mocap\"\n",
    "    opentrack_data_dir.parent.mkdir(parents=True, exist_ok=True)\n",
    "    \n",
    "    # Remove old symlink/directory if it exists\n",
    "    if opentrack_data_dir.exists() or opentrack_data_dir.is_symlink():\n",
    "        if opentrack_data_dir.is_symlink():\n",
    "            opentrack_data_dir.unlink()\n",
    "        else:\n",
    "            import shutil\n",
    "            shutil.rmtree(opentrack_data_dir)\n",
    "    \n",
    "    # Create symlink\n",
    "    opentrack_data_dir.symlink_to(DATASETS_DIR, target_is_directory=True)\n",
    "    print(f\"\\n✓ Created symlink: {opentrack_data_dir} -> {DATASETS_DIR}\")\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",
    "    # Check both MODELS_DIR and REPO_DIR/logs\n",
    "    search_dirs = [MODELS_DIR, REPO_DIR / \"logs\"]\n",
    "    \n",
    "    all_experiments = []\n",
    "    for logs_dir in search_dirs:\n",
    "        if not logs_dir.exists():\n",
    "            continue\n",
    "        \n",
    "        # Find all matching experiments\n",
    "        experiments = [\n",
    "            d for d in logs_dir.iterdir() \n",
    "            if d.is_dir() and exp_name_pattern in d.name\n",
    "        ]\n",
    "        all_experiments.extend(experiments)\n",
    "    \n",
    "    if not all_experiments:\n",
    "        return None\n",
    "    \n",
    "    # Sort by modification time and return latest\n",
    "    latest = sorted(all_experiments, key=lambda x: x.stat().st_mtime, reverse=True)[0]\n",
    "    return latest.name\n",
    "\n",
    "\n",
    "def find_generated_videos(output_dir=None):\n",
    "    \"\"\"\n",
    "    Find all generated video files\n",
    "    \"\"\"\n",
    "    if output_dir is None:\n",
    "        output_dir = VIDEOS_DIR\n",
    "    else:\n",
    "        output_dir = Path(output_dir)\n",
    "    \n",
    "    videos = []\n",
    "    \n",
    "    if output_dir.exists():\n",
    "        videos = list(output_dir.glob(\"*.mp4\")) + list(output_dir.glob(\"*.gif\"))\n",
    "    \n",
    "    if not videos:\n",
    "        # Try alternative locations\n",
    "        alternative_dirs = [REPO_DIR, REPO_DIR / \"logs\", MODELS_DIR]\n",
    "        for alt_dir in alternative_dirs:\n",
    "            if alt_dir.exists():\n",
    "                found = list(alt_dir.glob(\"**/*.mp4\")) + list(alt_dir.glob(\"**/*.gif\"))\n",
    "                videos.extend(found)\n",
    "    \n",
    "    return sorted(videos, key=lambda x: x.stat().st_mtime, reverse=True)\n",
    "\n",
    "\n",
    "def setup_model_output_symlink():\n",
    "    \"\"\"\n",
    "    Create symlink from OpenTrack's logs directory to our models directory\n",
    "    \"\"\"\n",
    "    opentrack_logs_dir = REPO_DIR / \"logs\"\n",
    "    \n",
    "    # Remove old symlink/directory if it exists\n",
    "    if opentrack_logs_dir.exists() or opentrack_logs_dir.is_symlink():\n",
    "        if opentrack_logs_dir.is_symlink():\n",
    "            opentrack_logs_dir.unlink()\n",
    "        else:\n",
    "            import shutil\n",
    "            # Move existing logs to MODELS_DIR first\n",
    "            if opentrack_logs_dir.is_dir():\n",
    "                for item in opentrack_logs_dir.iterdir():\n",
    "                    shutil.move(str(item), str(MODELS_DIR))\n",
    "            shutil.rmtree(opentrack_logs_dir)\n",
    "    \n",
    "    # Create symlink\n",
    "    opentrack_logs_dir.symlink_to(MODELS_DIR, target_is_directory=True)\n",
    "    print(f\"✓ Created symlink: {opentrack_logs_dir} -> {MODELS_DIR}\")\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": [
    "# Setup symlink for model outputs\n",
    "setup_model_output_symlink()\n",
    "\n",
    "print(f\"\\nModels will be saved to: {MODELS_DIR}\\n\")\n",
    "\n",
    "# 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",
    "    print(f\"✓ Model saved in: {MODELS_DIR / 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",
    "    print(f\"Videos will be saved to: {VIDEOS_DIR}\\n\")\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",
    "    # Move videos to VIDEOS_DIR if they're not already there\n",
    "    if videos:\n",
    "        import shutil\n",
    "        moved_videos = []\n",
    "        for video in videos:\n",
    "            if not video.is_relative_to(VIDEOS_DIR):\n",
    "                dest = VIDEOS_DIR / video.name\n",
    "                shutil.copy2(video, dest)\n",
    "                moved_videos.append(dest)\n",
    "                print(f\"Moved: {video.name} -> {dest}\")\n",
    "            else:\n",
    "                moved_videos.append(video)\n",
    "        videos = moved_videos\n",
    "    \n",
    "    if videos:\n",
    "        print(f\"\\n✓ Found {len(videos)} video(s) in {VIDEOS_DIR}:\")\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(REPO_DIR.rglob(\"*.mp4\")) + list(REPO_DIR.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 from our videos directory\n",
    "videos = find_generated_videos()\n",
    "\n",
    "if not videos:\n",
    "    # Try alternative search in the whole project\n",
    "    videos = list(REPO_DIR.rglob(\"*.mp4\")) + list(REPO_DIR.rglob(\"*.gif\"))\n",
    "\n",
    "if videos:\n",
    "    print(f\"Displaying {len(videos)} video(s) from {VIDEOS_DIR}:\\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\"Location: {video_path}\")\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",
    "**Project Structure:**\n",
    "```\n",
    "/data/workspaces/opentrack/\n",
    "├── datasets/          # Motion capture data (.npz files)\n",
    "│   └── lafan1/\n",
    "│       └── UnitreeG1/\n",
    "├── models/            # Trained models and checkpoints\n",
    "│   └── <timestamp>_<exp_name>/\n",
    "├── videos/            # Generated videos (.mp4, .gif)\n",
    "└── OpenTrack/         # Cloned repository\n",
    "    ├── data/mocap -> ../datasets/  (symlink)\n",
    "    └── logs -> ../models/          (symlink)\n",
    "```\n",
    "\n",
    "**Next Steps:**\n",
    "- All mocap data is in `/data/workspaces/opentrack/datasets/`\n",
    "- Run full training with GPU support\n",
    "- Test on rough terrain\n",
    "- Experiment with reference motion playback\n",
    "- Trained models are saved in `/data/workspaces/opentrack/models/`\n",
    "- Generated videos are in `/data/workspaces/opentrack/videos/`\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 `/data/workspaces/opentrack/models/` directory for experiment outputs\n",
    "- All data persists in `/data/workspaces/opentrack/` across notebook sessions"
   ]
  }
 ],
 "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
}