{
"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'![](\"{video_path}\")
'))\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",
"│ └── _/\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
}