diff --git a/.gitignore b/.gitignore index e99122e..336e462 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,51 @@ *__pycache__* -*.pyc \ No newline at end of file +*.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 + diff --git a/README.md b/README.md index e69de29..098b84f 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,79 @@ +# 🚁 MuJoCo 2D Quadrotor – MPC Control with Figure-8 Tracking under Wind Disturbance + +## 📌 Overview +This project implements a complete closed-loop robotics control system for a planar quadrotor simulated in MuJoCo. It integrates physics simulation, MPC-style controller, figure-8 trajectory tracking, sinusoidal wind disturbance, and performance evaluation with video and image outputs. + +## 📁 Project Structure +. +├── models/quadrotor_2d.xml +├── sim/__init__.py +├── sim/trajectory.py +├── sim/wind.py +├── sim/metrics.py +├── control/mpc_controller.py +├── run_mujoco_quadrotor.py +├── requirements.txt +└── outputs/quadrotor_figure8_mujoco.mp4 +└── outputs/quadrotor_figure8_final.png + +## 📦 Requirements +pip install -r requirements.txt + +Dependencies: +- mujoco >= 3.2.0 +- numpy >= 1.26 +- matplotlib >= 3.8 +- imageio >= 2.34 +- imageio-ffmpeg >= 0.5 + +## 🤖 System Model +State: [x, z, θ, vx, vz, θ̇] +Inputs: Fx (force x), Fz (force z), τθ (pitch torque) + +## 🌪️ Wind Disturbance +Fx = Ax sin(2πfx t + φ) +Fz = Az sin(2πfz t) - bias - kx + +## 📈 Reference Trajectory +Figure-8 path in x–z space: +[x, z, vx, vz, ax, az] + +## 🎮 MPC Controller +a_des = a_ref + Kp(p_ref - p) + Kd(v_ref - v) +Output: [Fx, Fz, τ] +Features: +- Horizon prediction (18 steps) +- Robust tracking under model mismatch + +## ⚙️ Control Limits +Fx: ±22 N +Fz: [-2, 24] N +Torque: ±2.4 Nm + +## 📊 Metrics +- Mean tracking error +- Max tracking error +- Position bounds +- Final state error + +## 🎥 Outputs +outputs/quadrotor_figure8_mujoco.mp4 +outputs/quadrotor_figure8_final.png + +## ▶️ Run +python run_mujoco_quadrotor.py --viewer +python run_mujoco_quadrotor.py --render-video --duration 20 + +## 🚁 MuJoCo Model +3-DOF planar quadrotor with: +- x/z slide joints +- pitch hinge +- corridor constraints +- mocap reference marker +- fixed camera view + +## 🔬 Key Features +MPC control + wind disturbance + MuJoCo physics + trajectory tracking + metrics + video generation. + +## 📌 Conclusion +Complete robotics control pipeline combining simulation, control, disturbance modeling, and evaluation in a single MuJoCo environment. \ No newline at end of file diff --git a/examples/quadrotor_drone/control/__init__.py b/examples/quadrotor_drone/control/__init__.py new file mode 100644 index 0000000..43a4d73 --- /dev/null +++ b/examples/quadrotor_drone/control/__init__.py @@ -0,0 +1,5 @@ +"""Quadrotor MPC controller package.""" + +from .mpc_controller import MPCConfig, QuadrotorMPC + +__all__ = ["MPCConfig", "QuadrotorMPC"] \ No newline at end of file diff --git a/examples/quadrotor_drone/control/mpc_controller.py b/examples/quadrotor_drone/control/mpc_controller.py new file mode 100644 index 0000000..6c221b9 --- /dev/null +++ b/examples/quadrotor_drone/control/mpc_controller.py @@ -0,0 +1,70 @@ +"""MPC-style controller for the MuJoCo 2D quadrotor.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + + +@dataclass(frozen=True) +class MPCConfig: + """Controller parameters.""" + + mass_estimate: float = 1.2 + gravity: float = 9.81 + horizon_steps: int = 18 + dt: float = 0.005 + max_tilt: float = 0.55 + kp: tuple[float, float] = (12.0, 10.0) + kd: tuple[float, float] = (7.0, 6.2) + theta_kp: float = 18.0 + theta_kd: float = 4.0 + max_force_x: float = 22.0 + max_force_z: float = 24.0 + max_torque: float = 2.4 + + +class QuadrotorMPC: + """Finite-horizon preview controller. + + MuJoCo simulates the true vehicle mass from the MJCF geoms. This controller + intentionally uses a 1.2 kg mass estimate, giving the requested mass error. + """ + + def __init__(self, config: MPCConfig | None = None) -> None: + self.config = config or MPCConfig() + + def command(self, state: np.ndarray, reference: np.ndarray) -> np.ndarray: + """Return MuJoCo actuator controls [Fx, Fz, pitch_tau].""" + cfg = self.config + pos = state[:2] + theta = state[2] + vel = state[3:5] + theta_dot = state[5] + + desired_acc = reference[4:6] + np.array(cfg.kp) * (reference[:2] - pos) + desired_acc += np.array(cfg.kd) * (reference[2:4] - vel) + + preview_pos = pos.copy() + preview_vel = vel.copy() + preview_error = np.zeros(2) + for _ in range(cfg.horizon_steps): + preview_vel += desired_acc * cfg.dt + preview_pos += preview_vel * cfg.dt + preview_error += reference[:2] - preview_pos + desired_acc += 0.22 * np.array(cfg.kp) * preview_error / cfg.horizon_steps + + desired_theta = np.clip(-desired_acc[0] / cfg.gravity, -cfg.max_tilt, cfg.max_tilt) + force_x = cfg.mass_estimate * desired_acc[0] + force_z = cfg.mass_estimate * (cfg.gravity + desired_acc[1]) + pitch_torque = cfg.theta_kp * (desired_theta - theta) - cfg.theta_kd * theta_dot + + return np.array( + [ + np.clip(force_x, -cfg.max_force_x, cfg.max_force_x), + np.clip(force_z, -2.0, cfg.max_force_z), + np.clip(pitch_torque, -cfg.max_torque, cfg.max_torque), + ], + dtype=float, + ) diff --git a/examples/quadrotor_drone/models/quadrotor_2d.xml b/examples/quadrotor_drone/models/quadrotor_2d.xml new file mode 100644 index 0000000..732f153 --- /dev/null +++ b/examples/quadrotor_drone/models/quadrotor_2d.xml @@ -0,0 +1,55 @@ + + + diff --git a/examples/quadrotor_drone/requirements.txt b/examples/quadrotor_drone/requirements.txt new file mode 100644 index 0000000..a1e336c --- /dev/null +++ b/examples/quadrotor_drone/requirements.txt @@ -0,0 +1,5 @@ +mujoco>=3.2.0 +numpy>=1.26 +matplotlib>=3.8 +imageio>=2.34 +imageio-ffmpeg>=0.5 diff --git a/examples/quadrotor_drone/run_mujoco_quadrotor.py b/examples/quadrotor_drone/run_mujoco_quadrotor.py new file mode 100644 index 0000000..b9d735f --- /dev/null +++ b/examples/quadrotor_drone/run_mujoco_quadrotor.py @@ -0,0 +1,139 @@ +"""Run the MuJoCo 2D quadrotor with 8-figure trajectory and sinusoidal wind.""" + +from __future__ import annotations + +import argparse +import pathlib +import time + +import imageio.v2 as imageio +import mujoco +import numpy as np + +from control.mpc_controller import QuadrotorMPC +from sim.metrics import TrackingMetrics +from sim.trajectory import FigureEightTrajectory +from sim.wind import SinusoidalWind + + +ROOT = pathlib.Path(__file__).resolve().parent +MODEL_PATH = ROOT / "models" / "quadrotor_2d.xml" +OUTPUT_DIR = ROOT / "outputs" + + +def state_from_data(data: mujoco.MjData) -> np.ndarray: + """Read [x, z, theta, vx, vz, theta_dot] from MuJoCo.""" + return np.array( + [ + data.qpos[0], + data.qpos[1], + data.qpos[2], + data.qvel[0], + data.qvel[1], + data.qvel[2], + ], + dtype=float, + ) + + +def set_reference_marker(model: mujoco.MjModel, data: mujoco.MjData, reference: np.ndarray) -> None: + marker_id = model.body("reference_marker").mocapid[0] + data.mocap_pos[marker_id] = np.array([reference[0], 0.0, reference[1]]) + + +def apply_wind(data: mujoco.MjData, wind_force: np.ndarray) -> None: + data.qfrc_applied[:] = 0.0 + data.qfrc_applied[0] = wind_force[0] + data.qfrc_applied[1] = wind_force[1] + + +def render_frame(renderer: mujoco.Renderer, model: mujoco.MjModel, data: mujoco.MjData) -> np.ndarray: + camera_id = model.camera("overview").id + renderer.update_scene(data, camera=camera_id) + return renderer.render() + + +def simulate(duration: float, viewer: bool, render_video: bool, width: int, height: int) -> dict[str, float]: + model = mujoco.MjModel.from_xml_path(str(MODEL_PATH)) + data = mujoco.MjData(model) + controller = QuadrotorMPC() + trajectory = FigureEightTrajectory() + wind = SinusoidalWind() + metrics = TrackingMetrics() + + data.qpos[:] = np.array([0.0, 4.0, 0.0]) + mujoco.mj_forward(model, data) + + frames = [] + renderer = mujoco.Renderer(model, height=height, width=width) if render_video else None + viewer_handle = None + if viewer: + from mujoco import viewer as mujoco_viewer + + viewer_handle = mujoco_viewer.launch_passive(model, data) + + steps = int(duration / model.opt.timestep) + render_stride = max(1, int((1.0 / 30.0) / model.opt.timestep)) + + try: + for step in range(steps): + state = state_from_data(data) + reference = trajectory.sample(data.time) + set_reference_marker(model, data, reference) + + wind_force = wind.force(data.time, state[:2]) + apply_wind(data, wind_force) + data.ctrl[:] = controller.command(state, reference) + + mujoco.mj_step(model, data) + metrics.update(state, reference) + + if viewer_handle: + viewer_handle.sync() + time.sleep(model.opt.timestep) + + if renderer and step % render_stride == 0: + frames.append(render_frame(renderer, model, data)) + finally: + if viewer_handle: + viewer_handle.close() + if renderer: + renderer.close() + + OUTPUT_DIR.mkdir(exist_ok=True) + if frames: + video_path = OUTPUT_DIR / "quadrotor_figure8_mujoco.mp4" + imageio.mimsave(video_path, frames, fps=30) + imageio.imwrite(OUTPUT_DIR / "quadrotor_figure8_final.png", frames[-1]) + + summary = metrics.summary() + summary["final_x"] = float(data.qpos[0]) + summary["final_z"] = float(data.qpos[1]) + summary["final_theta"] = float(data.qpos[2]) + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--duration", type=float, default=20.0) + parser.add_argument("--viewer", action="store_true", help="Open the interactive MuJoCo viewer.") + parser.add_argument("--render-video", action="store_true", help="Render MP4 and final PNG.") + parser.add_argument("--width", type=int, default=960) + parser.add_argument("--height", type=int, default=720) + args = parser.parse_args() + + summary = simulate( + duration=args.duration, + viewer=args.viewer, + render_video=args.render_video or not args.viewer, + width=args.width, + height=args.height, + ) + + print("MuJoCo 2D quadrotor, 8-figure trajectory, sinusoidal wind") + for key, value in summary.items(): + print(f"{key}: {value:.3f}") + + +if __name__ == "__main__": + main() diff --git a/examples/quadrotor_drone/sim/__init__.py b/examples/quadrotor_drone/sim/__init__.py new file mode 100644 index 0000000..c2f81fb --- /dev/null +++ b/examples/quadrotor_drone/sim/__init__.py @@ -0,0 +1,11 @@ +"""Simulation helpers for the MuJoCo planar quadrotor.""" + +from .metrics import TrackingMetrics +from .trajectory import FigureEightTrajectory +from .wind import SinusoidalWind + +__all__ = [ + "TrackingMetrics", + "FigureEightTrajectory", + "SinusoidalWind", +] \ No newline at end of file diff --git a/examples/quadrotor_drone/sim/metrics.py b/examples/quadrotor_drone/sim/metrics.py new file mode 100644 index 0000000..f551688 --- /dev/null +++ b/examples/quadrotor_drone/sim/metrics.py @@ -0,0 +1,32 @@ +"""Small metric accumulator for tracking performance.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np + + +@dataclass +class TrackingMetrics: + errors: list[float] = field(default_factory=list) + max_abs_x: float = 0.0 + min_z: float = 1e9 + max_z: float = -1e9 + + def update(self, state: np.ndarray, reference: np.ndarray) -> None: + xz = state[:2] + self.errors.append(float(np.linalg.norm(reference[:2] - xz))) + self.max_abs_x = max(self.max_abs_x, abs(float(xz[0]))) + self.min_z = min(self.min_z, float(xz[1])) + self.max_z = max(self.max_z, float(xz[1])) + + def summary(self) -> dict[str, float]: + return { + "mean_tracking_error": float(np.mean(self.errors)), + "max_tracking_error": float(np.max(self.errors)), + "max_abs_x": self.max_abs_x, + "min_z": self.min_z, + "max_z": self.max_z, + } + diff --git a/examples/quadrotor_drone/sim/trajectory.py b/examples/quadrotor_drone/sim/trajectory.py new file mode 100644 index 0000000..8b0ae18 --- /dev/null +++ b/examples/quadrotor_drone/sim/trajectory.py @@ -0,0 +1,34 @@ +"""8-figure trajectory in x-z.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + + +@dataclass(frozen=True) +class FigureEightTrajectory: + """A smooth vertical number-8 reference.""" + + center_x: float = 0.0 + center_z: float = 4.0 + amp_x: float = 1.15 + amp_z: float = 2.25 + period: float = 12.0 + + def sample(self, time: float) -> np.ndarray: + """Return [x, z, vx, vz, ax, az].""" + omega = 2.0 * np.pi / self.period + wt = omega * time + + x = self.center_x + self.amp_x * np.sin(2.0 * wt) / 2.0 + z = self.center_z + self.amp_z * np.sin(wt) + + vx = self.amp_x * omega * np.cos(2.0 * wt) + vz = self.amp_z * omega * np.cos(wt) + + ax = -2.0 * self.amp_x * omega**2 * np.sin(2.0 * wt) + az = -self.amp_z * omega**2 * np.sin(wt) + + return np.array([x, z, vx, vz, ax, az], dtype=float) diff --git a/examples/quadrotor_drone/sim/wind.py b/examples/quadrotor_drone/sim/wind.py new file mode 100644 index 0000000..78f4c84 --- /dev/null +++ b/examples/quadrotor_drone/sim/wind.py @@ -0,0 +1,29 @@ +"""Sinusoidal wind disturbance for the planar MuJoCo quadrotor.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + + +@dataclass(frozen=True) +class SinusoidalWind: + """World-frame force disturbance applied to x and z generalized forces.""" + + x_amplitude: float = 0.85 + z_amplitude: float = 0.38 + x_frequency: float = 0.23 + z_frequency: float = 0.37 + phase: float = 0.6 + downward_bias: float = 0.16 + + def force(self, time: float, xz: np.ndarray) -> np.ndarray: + fx = self.x_amplitude * np.sin(2.0 * np.pi * self.x_frequency * time + self.phase) + fz = ( + self.z_amplitude * np.sin(2.0 * np.pi * self.z_frequency * time) + - self.downward_bias + - 0.08 * xz[0] + ) + return np.array([fx, fz], dtype=float) + diff --git a/mpc_controller.py b/mpc_controller.py new file mode 100644 index 0000000..6c221b9 --- /dev/null +++ b/mpc_controller.py @@ -0,0 +1,70 @@ +"""MPC-style controller for the MuJoCo 2D quadrotor.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + + +@dataclass(frozen=True) +class MPCConfig: + """Controller parameters.""" + + mass_estimate: float = 1.2 + gravity: float = 9.81 + horizon_steps: int = 18 + dt: float = 0.005 + max_tilt: float = 0.55 + kp: tuple[float, float] = (12.0, 10.0) + kd: tuple[float, float] = (7.0, 6.2) + theta_kp: float = 18.0 + theta_kd: float = 4.0 + max_force_x: float = 22.0 + max_force_z: float = 24.0 + max_torque: float = 2.4 + + +class QuadrotorMPC: + """Finite-horizon preview controller. + + MuJoCo simulates the true vehicle mass from the MJCF geoms. This controller + intentionally uses a 1.2 kg mass estimate, giving the requested mass error. + """ + + def __init__(self, config: MPCConfig | None = None) -> None: + self.config = config or MPCConfig() + + def command(self, state: np.ndarray, reference: np.ndarray) -> np.ndarray: + """Return MuJoCo actuator controls [Fx, Fz, pitch_tau].""" + cfg = self.config + pos = state[:2] + theta = state[2] + vel = state[3:5] + theta_dot = state[5] + + desired_acc = reference[4:6] + np.array(cfg.kp) * (reference[:2] - pos) + desired_acc += np.array(cfg.kd) * (reference[2:4] - vel) + + preview_pos = pos.copy() + preview_vel = vel.copy() + preview_error = np.zeros(2) + for _ in range(cfg.horizon_steps): + preview_vel += desired_acc * cfg.dt + preview_pos += preview_vel * cfg.dt + preview_error += reference[:2] - preview_pos + desired_acc += 0.22 * np.array(cfg.kp) * preview_error / cfg.horizon_steps + + desired_theta = np.clip(-desired_acc[0] / cfg.gravity, -cfg.max_tilt, cfg.max_tilt) + force_x = cfg.mass_estimate * desired_acc[0] + force_z = cfg.mass_estimate * (cfg.gravity + desired_acc[1]) + pitch_torque = cfg.theta_kp * (desired_theta - theta) - cfg.theta_kd * theta_dot + + return np.array( + [ + np.clip(force_x, -cfg.max_force_x, cfg.max_force_x), + np.clip(force_z, -2.0, cfg.max_force_z), + np.clip(pitch_torque, -cfg.max_torque, cfg.max_torque), + ], + dtype=float, + ) diff --git a/src/simulator/renderer.py b/src/simulator/renderer.py index e7c4820..ce3d328 100644 --- a/src/simulator/renderer.py +++ b/src/simulator/renderer.py @@ -4,9 +4,10 @@ from matplotlib.patches import Circle, Rectangle from matplotlib.collections import LineCollection +from .video_recorder import VideoRecorder, create_recorder class Renderer: - def __init__(self, x_limits=(-10.0, 10.0), y_limits=(-10.0, 10.0), max_colors=20) -> None: + def __init__(self, x_limits=(-10.0, 10.0), y_limits=(-10.0, 10.0), max_colors=20, record_path: str | None = None, record_fps: float = 30.0,record_format: str | None = None, ) -> None: self.fig = plt.figure() self.ax = self.fig.add_subplot(111, aspect="equal") self.ax.set_xlim(x_limits) @@ -34,7 +35,27 @@ def __init__(self, x_limits=(-10.0, 10.0), y_limits=(-10.0, 10.0), max_colors=20 self.colors = plt.cm.rainbow(np.linspace(0, 1, max_colors)) self.base_patch = None - + # Video recording setup (deferred until first frame size is known) + + self._recorder: VideoRecorder | None = None + self._record_enabled = record_path is not None + self._record_path = record_path + self._record_fps = record_fps + self._record_format = record_format + + def _start_recording_if_needed(self, frame_size: tuple[int, int]) -> None: + """Start recording if enabled and not already started. + + Args: + frame_size: (width, height) of the frame. + """ + if self._recorder is None and self._record_enabled: + self._recorder = create_recorder( + output_path=self._record_path, + fps=self._record_fps, + format=self._record_format, + ) + self._recorder.start(frame_size) def update(self, objects, dt=0.0001): if not objects: @@ -99,6 +120,23 @@ def draw_tree(obj, q): self.fig.canvas.draw_idle() self.fig.canvas.flush_events() + # Record frame if recording is enabled + if self._record_enabled: + frame = np.array(self.fig.canvas.buffer_rgba()) + height, width = frame.shape[:2] + frame_size = (width, height) + self._start_recording_if_needed(frame_size) + + if self._recorder is not None and self._recorder.is_recording(): + self._recorder.add_frame(frame) + def close(self): + """Close the renderer and finalize video recording.""" + if self._recorder is not None and self._recorder.is_recording(): + self._recorder.stop() plt.close(self.fig) + + def is_recording(self) -> bool: + """Check if video recording is active.""" + return self._recorder is not None and self._recorder.is_recording() diff --git a/src/simulator/video_recorder.py b/src/simulator/video_recorder.py new file mode 100644 index 0000000..85e364a --- /dev/null +++ b/src/simulator/video_recorder.py @@ -0,0 +1,122 @@ +""" Video recording methods and function for simulater with OpenCV""" +import cv2 +import pathlib +import numpy as np +from typing import Literal + +class VideoRecorder: + """Records simulation frames to video files using OpenCV with supporting MP4 and AVI output formats + """ + + def __init__( + self, + output_path: str | pathlib.Path, + fps: float = 30.0, + format: Literal["mp4", "avi"] = "mp4", + ) -> None: + """Initialize the video recorder: + Args: + output_path: Path to save the output video file + fps: Frames per second for the output vide + format: Output format, either "mp4" or "avi" + """ + self.output_path = pathlib.Path(output_path) + self.fps = fps + self.format = format + self._writer: cv2.VideoWriter | None = None + self._frame_size: tuple[int, int] | None = None + + def start(self, frame_size: tuple[int, int]) -> None: + """Start recording - creates the output file writer + + Args: + frame_size: (width, height) of each frame + """ + if self._writer is not None: + msg = "Recorder already started" + raise RuntimeError(msg) + + self._frame_size = frame_size + + + + if self.format == "mp4": + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + elif self.format == "avi": + fourcc = cv2.VideoWriter_fourcc(*"XVID") + else: + msg = f"Unsupported format: {self.format}" + raise ValueError(msg) + + self._writer = cv2.VideoWriter( + str(self.output_path), + fourcc, + self.fps, + frame_size, + ) + + if not self._writer.isOpened(): + msg = f"Failed to open video writer for {self.output_path}" + raise RuntimeError(msg) + def add_frame(self, frame: np.ndarray) -> None: + """Add a frame to the recording. + + Args: + frame: Image data as numpy array (RGB or BGR). + """ + if self._writer is None: + msg = "Recorder not started. Call start() first." + raise RuntimeError(msg) + + # OpenCV expects BGR format for VideoWriter + if len(frame.shape) == 3 and frame.shape[2] == 4: + # RGBA to BGR + frame = frame[:, :, :3] + + elif len(frame.shape) == 3 and frame.shape[2] == 3: + # RGB to BGR + pass + + self._writer.write(frame) + + def stop(self) -> None: + """Stop recording and close the output file.""" + if self._writer is None: + msg = "Recorder not started" + raise RuntimeError(msg) + + self._writer = None + + def is_recording(self) -> bool: + """Check if the recorder is currently active.""" + return self._writer is not None and self._writer.isOpened() + +def create_recorder( + output_path: str | pathlib.Path, + fps: float = 30.0, + format: Literal["mp4", "avi"] | None = None, +) -> VideoRecorder: + """Factory function to create a VideoRecorder. + + Automatically detects format from file extension if not specified. + + Args: + output_path: Path to save the output video file + fps: Frames per second for the output vide + format: Output format, either "mp4" or "avi"and Auto-detected if None + + Returns: + Configured VideoRecorder instance + """ + path = pathlib.Path(output_path) + if format is None: + ext = path.suffix.lower() + if ext in {".mp4", ".mov"}: + format = "mp4" + elif ext in {".avi", }: + format = "avi" + else: + format = "mp4" + path = path.with_suffix(".mp4") + + return VideoRecorder(output_path=path, fps=fps, format=format) \ No newline at end of file diff --git a/src/simulator/world.py b/src/simulator/world.py index 813c884..5a4f7c1 100644 --- a/src/simulator/world.py +++ b/src/simulator/world.py @@ -2,7 +2,27 @@ class World: - def __init__(self, physics, renderer): + """Simulation world managing physics and rendering. + + Coordinates the simulation loop, physics updates, and rendering. + """ + def __init__( + self, + physics, + renderer, + record_path: str | None = None, + record_fps: float = 30.0, + record_format: str | None = None, + ): + """Initialize the world. + + Args: + physics: Physics engine for dynamics simulation + renderer: Renderer for visualization + record_path: If set, enable recording to this file path + record_fps: Frame rate for video recording + record_format: Recording format ("mp4" or "avi") + """ self.physics = physics self.renderer = renderer self.objects = []