Upload opentrack.ipynb

samples/opentrack.ipynb

@@ -33,7 +33,19 @@
     "import threading\n",
     "import queue\n",
     "\n",
-    "print(\"Environment setup complete!\")"
+    "# 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!\")"
    ]
   },
   {
@@ -50,14 +62,19 @@
    "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",
+    "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(\"✓ Repository already exists\")\n",
+    "    print(f\"✓ Repository already exists at {REPO_DIR}\")\n",
     "\n",
-    "os.chdir('OpenTrack')\n",
-    "print(f\"Current directory: {os.getcwd()}\")"
+    "# Change to repository directory\n",
+    "os.chdir(REPO_DIR)\n",
+    "print(f\"\\n✓ Working directory: {os.getcwd()}\")"
    ]
   },
   {
@@ -100,23 +117,23 @@
    "source": [
     "from huggingface_hub import snapshot_download\n",
     "\n",
-    "# Create directory structure\n",
-    "mocap_dir = Path(\"data/mocap/lafan1/UnitreeG1\")\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",
-    "    # 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=str(DATASETS_DIR),\n",
     "        local_dir_use_symlinks=False\n",
     "    )\n",
     "    \n",
@@ -133,6 +150,22 @@
     "        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",
@@ -190,38 +223,76 @@
     "    \"\"\"\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",
+    "    # Check both MODELS_DIR and REPO_DIR/logs\n",
+    "    search_dirs = [MODELS_DIR, REPO_DIR / \"logs\"]\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",
+    "    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",
-    "    return experiments[0].name if experiments else None\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=\"videos\"):\n",
+    "def find_generated_videos(output_dir=None):\n",
     "    \"\"\"\n",
     "    Find all generated video files\n",
     "    \"\"\"\n",
-    "    video_dir = Path(output_dir)\n",
-    "    if not video_dir.exists():\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 = [\".\", \"logs\", \"outputs\"]\n",
+    "        alternative_dirs = [REPO_DIR, REPO_DIR / \"logs\", MODELS_DIR]\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",
+    "            if alt_dir.exists():\n",
+    "                found = list(alt_dir.glob(\"**/*.mp4\")) + list(alt_dir.glob(\"**/*.gif\"))\n",
+    "                videos.extend(found)\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",
+    "\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\")"
    ]
   },
@@ -240,6 +311,11 @@
    "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",
@@ -253,6 +329,7 @@
     "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\")"
    ]
@@ -306,6 +383,7 @@
     "\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",
@@ -327,14 +405,28 @@
     "    # 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):\")\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(Path(\".\").rglob(\"*.mp4\")) + list(Path(\".\").rglob(\"*.gif\"))\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",
@@ -356,19 +448,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Find and display all generated videos\n",
+    "# Find and display all generated videos from our videos directory\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",
+    "    # 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):\\n\")\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",
@@ -456,16 +549,33 @@
     "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",
-    "- Download more mocap files for better training\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 `logs/` directory for experiment outputs"
+    "- Check `/data/workspaces/opentrack/models/` directory for experiment outputs\n",
+    "- All data persists in `/data/workspaces/opentrack/` across notebook sessions"
    ]
   }
  ],