Renderer add recording / video export #25

.gitignore

@@ -1,2 +1,51 @@
 *__pycache__*
-*.pyc
+*.pyc
+# VS Code files
+.vscode/
+*.code-workspace
+.vscode-test/
+
+# Python environment
+venv/
+env/
+ENV/
+.venv/
+virtualenv/
+env.bak/
+venv.bak/
+
+# Python virtual environment files
+lib/
+include/
+bin/
+share/
+pyvenv.cfg
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Log files
+*.log
+logs/
+*.log.*
+log.txt
+debug.log
+error.log
+
+# OS generated files
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db
+*.swp
+*.swo
+
+# Python cache
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+

README.md

@@ -0,0 +1,149 @@
+# Robot Simulator
+
+A Python-based robot arm simulator with video recording capability.
+
+
+
+## Problem: How to Record Robot Simulations?
+
+### The Challenge
+
+We needed to add video recording capability to a matplotlib-based robot simulator. The requirement was to:
+1. Capture each frame of the simulation
+2. Save it as a video file (MP4/AVI)
+3. Use OpenCV for future extensibility (robotics + AI projects)
+
+### Our Solution
+
+We used **OpenCV** (`cv2.VideoWriter`) to capture and encode video frames:
+
+1. **Frame Capture**: After matplotlib renders each frame, we capture it from the canvas buffer
+2. **Color Conversion**: matplotlib gives RGBA format, OpenCV needs BGR - we convert using `cv2.cvtColor()`
+3. **Video Encoding**: OpenCV's `VideoWriter` handles MP4/AVI encoding efficiently
+4. **Auto Folder**: Recordings automatically save to `recording_simulation/` folder
+
+### Why OpenCV Instead of imageio?
+
+| Feature | OpenCV | imageio |
+|---------|--------|---------|
+| Computer Vision | ✅ Yes | ❌ No |
+| Real-time camera support | ✅ Yes | ❌ No |
+| ML/AI integration | ✅ Yes | ❌ No |
+| Industry standard for robotics | ✅ Yes | ❌ No |
+
+OpenCV is the better choice for robotics + AI projects since you may later want to:
+- Process video with ML models
+- Add real camera feeds
+- Do object detection (YOLO, etc.)
+
+## Files Added/Modified
+
+### New Files
+- **`src/simulator/video_recorder.py`** - OpenCV-based video recording class
+- **`src/simulator/spatial/__init__.py`** - Module initialization for relative imports
+
+### Modified Files
+- **`src/simulator/renderer.py`** - Added frame capture and recording to `update()` method
+- **`src/simulator/main.py`** - Added CLI arguments for recording control
+- **`src/simulator/world.py`** - Simplified recording setup
+- **`src/simulator/dynamics/ab_algorithm.py`** - Fixed scalar extraction bug
+- **`pyproject.toml`** - Added `opencv-python` dependency
+
+## How to Run
+
+### Prerequisites
+
+Install dependencies:
+```bash
+pip install opencv-python matplotlib scipy numpy
+```
+
+### Basic Usage
+
+```bash
+# Run from the simulator-master folder
+cd simulator-master
+
+# Run simulation with default robot (CartPole)
+python src/simulator/main.py --steps 300
+
+# Record to video (saved in recording_simulation/ folder)
+python src/simulator/main.py --record my_robot.mp4 --steps 200 --fps 15
+```
+
+### Command Line Options
+
+| Option | Description | Default |
+|--------|-------------|---------|
+| `--robot` | Robot type: `cartpole`, `two-link`, `tree7`, `robot-tree` | `cartpole` |
+| `--steps` | Number of simulation steps | `1000` |
+| `--record` | Output filename for recording | (no recording) |
+| `--fps` | Frames per second for video | `30.0` |
+| `--format` | Video format: `mp4`, `avi` | (auto-detected) |
+
+### Examples
+
+```bash
+# Record CartPole (default robot)
+python src/simulator/main.py --record cartpole_demo.mp4 --steps 300 --fps 20
+
+# Record Two-Link robot
+python src/simulator/main.py --robot two-link --record two_link.mp4 --steps 200
+
+# Record Tree7 robot with custom FPS
+python src/simulator/main.py --robot tree7 --record tree7_sim.mp4 --fps 30
+
+# Record with custom path
+python src/simulator/main.py --record /path/to/my_video.mp4 --steps 150
+```
+
+## Output Location
+
+Videos are saved in the `recording_simulation/` folder (automatically created):
+
+```
+TEAM-005/
+├── recording_simulation/
+│   ├── my_robot.mp4
+│   ├── cartpole_demo.mp4
+│   └── tree7_sim.mp4
+└── ...
+```
+
+## Robot Models Available
+
+1. **CartPole** - Simple pole on a cart (default)
+2. **TwoLink** - Two-segment arm
+3. **Tree7** - 7-DOF tree structure
+4. **RobotTree** - Custom tree with configurable DOF
+
+## Technical Details
+
+### Dynamics Computation
+- Uses **Articulated Body Algorithm (ABA)** from Featherstone's Rigid Body Dynamics
+- Integrates equations of motion using `scipy.integrate.solve_ivp`
+- Supports gravity, joint torques, and external forces
+
+### Rendering
+- Matplotlib-based 2D visualization
+- LineCollection for efficient link drawing
+- Scattered points for joints
+- World frame axes displayed (red=X, green=Y)
+
+### Recording Process
+1. Each simulation step renders the robot
+2. `renderer.update()` captures the canvas as RGBA numpy array
+3. `VideoRecorder.add_frame()` converts RGBA→BGR and writes to file
+4. On completion, `VideoRecorder.stop()` finalizes the video file
+
+## Future Extensions
+
+With OpenCV installed, you can now:
+- Add real camera feeds to the simulation
+- Implement ML-based controllers
+- Use computer vision for feedback control
+- Integrate object detection for target tracking
+- Add image processing pipelines
+
+
+

src/simulator/dynamics/ab_algorithm.py

@@ -1,7 +1,7 @@
 import numpy as np
 
-from spatial.so import SO3, so3
-from spatial.se import SE3, se3, crm, crf
+from ..spatial.so import SO3, so3
+from ..spatial.se import SE3, se3, crm, crf
 
 
 class ABAlgorithm:
@@ -45,8 +45,8 @@ def forward_dynamics(self, obj, q, qd, tau, f_ext=None):
 
         for i in reversed(range(NB)):
             U[i] = IA[i] @ S[i]
-            d[i] = S[i].T @ U[i]
-            u[i] = tau[i] - S[i].T @ pA[i]
+            d[i] = (S[i].T @ U[i])[0, 0]  # Extract scalar from 1x1 matrix
+            u[i] = tau[i] - (S[i].T @ pA[i])[0, 0]
 
             parent = model["parent"][i]
             if parent != -1:
@@ -65,7 +65,7 @@ def forward_dynamics(self, obj, q, qd, tau, f_ext=None):
                 a[i] = Xup[i] @ (-a_grav) + c[i]
             else:
                 a[i] = Xup[i] @ a[parent] + c[i]
-            qdd[i] = (u[i] - U[i].T @ a[i]) / d[i]
+            qdd[i] = (u[i] - (U[i].T @ a[i])[0, 0]) / d[i]  # Extract scalar from 1x1 matrix
             a[i] = a[i] + S[i].dot(qdd[i])
 
         return qdd

src/simulator/main.py

@@ -1,34 +1,124 @@
+"""Main entry point for the robot simulator."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import pathlib
 import sys
-import numpy as np
 
-from physics import PhysicsEngine
-from renderer import Renderer
-from world import World
+# Add parent directory to path so we can import simulator as a package
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from simulator.dynamics.ab_algorithm import ABAlgorithm
+from simulator.objects import CartPole, RobotTree, Tree7, TwoLink
+from simulator.physics import PhysicsEngine
+from simulator.renderer import Renderer
+from simulator.world import World
+
+# Default recording folder
+DEFAULT_RECORDING_FOLDER = "recording_simulation"
+
+
+def resolve_record_path(record_arg: str | None) -> str | None:
+    """Resolve the recording path.
+
+    If record_arg is a filename (no directory separators), save to DEFAULT_RECORDING_FOLDER.
+    Otherwise use the provided path directly.
 
-from objects import RobotTree, TwoLink, Tree7, CartPole
-from dynamics.ab_algorithm import ABAlgorithm
+    Args:
+        record_arg: The --record argument value.
+
+    Returns:
+        Resolved path for recording, or None if recording is disabled.
+    """
+    if record_arg is None:
+        return None
+
+    record_path = pathlib.Path(record_arg)
+
+    # If it's just a filename (no directory), use the default recording folder
+    if record_path.parent == pathlib.Path("."):
+        recording_folder = pathlib.Path(DEFAULT_RECORDING_FOLDER)
+        recording_folder.mkdir(exist_ok=True)
+        record_path = recording_folder / record_path
+
+    return str(record_path)
 
 
 def main():
+    """Run the robot simulator."""
+    parser = argparse.ArgumentParser(
+        description="Robot arm simulator with video recording support"
+    )
+    parser.add_argument(
+        "--robot",
+        type=str,
+        default="cartpole",
+        choices=["two-link", "tree7", "robot-tree", "cartpole"],
+        help="Robot type to simulate (default: cartpole)",
+    )
+    parser.add_argument(
+        "--steps",
+        type=int,
+        default=1000,
+        help="Number of simulation steps (default: 1000)",
+    )
+    parser.add_argument(
+        "--record",
+        type=str,
+        metavar="FILENAME",
+        help="Record simulation to video (saved to recording_simulation/ folder)",
+    )
+    parser.add_argument(
+        "--fps",
+        type=float,
+        default=30.0,
+        help="Frame rate for video recording (default: 30.0)",
+    )
+    parser.add_argument(
+        "--format",
+        type=str,
+        choices=["mp4", "avi"],
+        help="Video format (auto-detected from extension if not specified)",
+    )
+    args = parser.parse_args()
+
+    # Resolve recording path
+    record_path = resolve_record_path(args.record)
 
     fd_solver = ABAlgorithm()
 
     physics = PhysicsEngine(fd_solver, gravity=[0.0, -9.81, 0.0])
-    renderer = Renderer()
+    renderer = Renderer(
+        record_path=record_path,
+        record_fps=args.fps,
+        record_format=args.format,
+    )
 
     world = World(physics, renderer)
 
-    # robot = TwoLink()
-    # robot = Tree7()
+    # Select robot based on argument
+    if args.robot == "two-link":
+        robot = TwoLink()
+    elif args.robot == "tree7":
+        robot = Tree7()
+    elif args.robot == "robot-tree":
+        robot = RobotTree()
+        robot.some_tree(8, 2)
+    else:  # cartpole (default)
+        robot = CartPole()
 
-    # robot = RobotTree()
-    # robot.some_tree(8,2)
+    world.add_object(robot)
 
-    robot = CartPole()
+    print(f"Starting simulation: {args.robot}")
+    if record_path:
+        print(f"Recording to: {record_path}")
 
-    world.add_object(robot)
+    world.run(args.steps)
 
-    world.run(1000)
+    if record_path:
+        print(f"Recording saved to: {record_path}")
 
 
 if __name__ == "__main__":