feat(rl): Gymnasium QuadcopterEnv with modular action/obs/reward/task system

src/drones_sim/rl/__init__.py

@@ -0,0 +1,11 @@
+"""Gymnasium environment for quadcopter reinforcement learning.
+
+The primary entry point is ``QuadcopterEnv``, which wraps
+``QuadcopterDynamics`` into a standard ``gymnasium.Env``.
+"""
+
+from .actions import MotorSpeedAction, ThrustBodyRatesAction  # noqa: F401
+from .env import QuadcopterEnv  # noqa: F401
+from .observations import RelativeStateObs  # noqa: F401
+from .reward import RewardConfig, reward  # noqa: F401
+from .tasks import HoverTask, TrackingTask, WaypointTask  # noqa: F401

src/drones_sim/rl/actions.py

@@ -0,0 +1,78 @@
+"""Action parameterizations for the quadcopter RL environment.
+
+Each parameterization defines ``low``, ``high``, ``dim`` and implements
+``to_motors(action) -> (4,)`` — converting a policy action vector into
+four motor speed commands (rad/s).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+class MotorSpeedAction:
+    """Lowest level: raw motor speeds for each of 4 rotors.
+
+    Hardest to learn, suitable for acrobatic maneuvers.
+    """
+
+    dim = 4
+    low = np.zeros(4, dtype=np.float32)
+    high = np.full(4, 4000.0, dtype=np.float32)
+
+    def to_motors(self, quad, action: NDArray) -> NDArray:
+        return np.clip(action, self.low, self.high)
+
+
+class ThrustBodyRatesAction:
+    """Mid-level: collective thrust + 3 body-rate commands.
+
+    Action: [thrust (N), omega_x, omega_y, omega_z (rad/s)].
+
+    The thrust is allocated equally across 4 motors (hover thrust added),
+    and the body-rate error is mapped to differential motor commands via
+    the quadcopter allocation matrix inverted as a proportional controller.
+    """
+
+    dim = 4
+    _thrust_low = 0.2 * 9.81
+    _thrust_high = 2.0 * 9.81
+    _rate_low = -5.0
+    _rate_high = 5.0
+
+    @property
+    def low(self) -> np.ndarray:
+        return np.array([self._thrust_low, self._rate_low, self._rate_low, self._rate_low], dtype=np.float32)
+
+    @property
+    def high(self) -> np.ndarray:
+        return np.array([self._thrust_high, self._rate_high, self._rate_high, self._rate_high], dtype=np.float32)
+
+    def to_motors(self, quad, action: NDArray) -> NDArray:
+        thrust, wx, wy, wz = action
+        # Hover feedforward
+        hover_thrust = quad.mass * quad.g
+        total_thrust = thrust + hover_thrust
+
+        # Rate errors → torque commands (proportional rate controller)
+        omega = quad.get_angular_velocity()
+        tau = np.array([wx - omega[0], wy - omega[1], wz - omega[2]])
+
+        # Build wrench and invert allocation
+        wrench = np.array([
+            np.clip(total_thrust, 0.2 * hover_thrust, 2.0 * hover_thrust),
+            tau[0] * 0.05,
+            tau[1] * 0.05,
+            tau[2] * 0.02,
+        ])
+
+        alloc = quad.allocation_matrix()
+        try:
+            w_sq = np.linalg.solve(alloc, wrench)
+        except np.linalg.LinAlgError:
+            hover_w = np.sqrt(hover_thrust / (4 * quad.k_f))
+            return np.full(4, hover_w)
+
+        motor_speeds = np.sqrt(np.maximum(w_sq, 0.0))
+        return np.clip(motor_speeds, 0.0, 4000.0)

src/drones_sim/rl/env.py

@@ -0,0 +1,189 @@
+"""Gymnasium environment wrapping QuadcopterDynamics.
+
+``QuadcopterEnv`` is the single integration point::
+
+    env = QuadcopterEnv(task=HoverTask(), action_param=ThrustBodyRatesAction(),
+                         obs_builder=RelativeStateObs(), reward_fn=reward_hover)
+    obs, _ = env.reset()
+    for _ in range(1000):
+        action = policy(obs)
+        obs, reward, terminated, truncated, info = env.step(action)
+
+It follows the standard ``gymnasium.Env`` contract and is compatible with
+Stable-Baselines3, CleanRL, Tianshou, and RLlib.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from drones_sim.dynamics import QuadcopterDynamics
+
+# gymnasium is an optional dependency (rl extra).  Every import site tests
+# for availability so the module can be imported for type-checking and
+# contract tests even without torch/gym installed.
+try:
+    import gymnasium as gym
+    from gymnasium import spaces
+except ImportError:  # pragma: no cover
+    gym = None  # type: ignore[assignment]
+    spaces = None  # type: ignore[assignment]
+
+from .actions import ThrustBodyRatesAction
+from .observations import RelativeStateObs
+from .reward import RewardConfig, reward
+from .tasks import HoverTask
+
+_GRAVITY = np.array([0.0, 0.0, 9.81])
+_TILT_LIMIT = np.deg2rad(75.0)  # crash threshold
+_POS_LIMIT = 50.0               # out-of-bounds
+
+
+class QuadcopterEnv(gym.Env if gym else object):
+    """Gymnasium wrapper around ``QuadcopterDynamics``.
+
+    Parameters
+    ----------
+    task : Task
+        Defines the target position/velocity (HoverTask, WaypointTask, etc.).
+    action_param : ActionParameterization
+        Converts policy action → motor speeds.
+    obs_builder : ObservationBuilder
+        Builds the observation vector from the quad state.
+    reward_fn : callable
+        ``reward(quad, task, action, step_idx, cfg, prev_action) -> float``.
+    reward_cfg : RewardConfig
+        Weights for the reward function.
+    dt : float
+        Simulation time-step [s].
+    episode_len_s : float
+        Maximum episode duration [s].
+    render_mode : str or None
+        'viser' for live 3D viewer (requires running viser server).
+    seed : int or None
+        Random seed for reproducibility.
+    """
+
+    metadata = {"render_modes": ["human", "viser"]}
+
+    def __init__(
+        self,
+        task=None,
+        action_param=None,
+        obs_builder=None,
+        reward_fn=None,
+        reward_cfg=None,
+        dt: float = 0.01,
+        episode_len_s: float = 10.0,
+        render_mode: str | None = None,
+        seed: int | None = None,
+    ):
+        if gym is None:
+            raise ImportError("gymnasium is required for QuadcopterEnv.  "
+                              "Install it with: pip install gymnasium")
+        super().__init__()
+        self.dt = dt
+        self.max_steps = int(episode_len_s / dt)
+        self.task = task if task is not None else HoverTask()
+        self.action_param = action_param if action_param is not None else ThrustBodyRatesAction()
+        self.obs_builder = obs_builder if obs_builder is not None else RelativeStateObs()
+        self.reward_fn = reward_fn if reward_fn is not None else reward
+        self.reward_cfg = reward_cfg if reward_cfg is not None else RewardConfig()
+        self.render_mode = render_mode
+
+        self.observation_space = spaces.Box(
+            low=-np.inf, high=np.inf,
+            shape=(self.obs_builder.dim,), dtype=np.float32,
+        )
+        self.action_space = spaces.Box(
+            low=self.action_param.low, high=self.action_param.high, dtype=np.float32,
+        )
+
+        self.quad = QuadcopterDynamics(motor_time_constant=0.04)
+        self._step_idx = 0
+        self._rng = np.random.default_rng(seed)
+        self._prev_action: np.ndarray | None = None
+
+    # ------------------------------------------------------------------
+    # gym.Env contract
+    # ------------------------------------------------------------------
+
+    def reset(self, *, seed: int | None = None, options: dict | None = None):
+        super().reset(seed=seed)
+        self._rng = np.random.default_rng(seed)
+        # Seed the global RNG for SensorNoiseModel compatibility
+        np.random.seed(seed)
+        self.quad.reset()
+        if hasattr(self.task, 'reset'):
+            self.task.reset(self.quad, rng=self._rng)
+        if hasattr(self.obs_builder, '_prev_action'):
+            self.obs_builder._prev_action = np.zeros(4, dtype=np.float32)
+        self._step_idx = 0
+        self._prev_action = None
+        # Pre-fill motor lag with hover-speed so the drone doesn't dip
+        # below z=0 on the very first step.
+        hover_w = np.sqrt(self.quad.mass * self.quad.g / (4 * self.quad.k_f))
+        self.quad.motor_states = np.full(4, hover_w)
+        obs = self.obs_builder.build(self.quad, self.task, action=None)
+        return obs.astype(np.float32), {"t": 0.0}
+
+    def step(self, action):
+        motor_speeds = self.action_param.to_motors(self.quad, action)
+        self.quad.update(self.dt, motor_speeds)
+        self._step_idx += 1
+
+        obs = self.obs_builder.build(self.quad, self.task, action=action)
+        r = self.reward_fn(
+            self.quad, self.task, action, self._step_idx,
+            self.reward_cfg, self._prev_action,
+        )
+        self._prev_action = action.copy()
+
+        terminated = self._is_crashed()
+        truncated = self._step_idx >= self.max_steps
+        info = {"t": self._step_idx * self.dt, "motor_speeds": motor_speeds}
+        return obs.astype(np.float32), float(r), terminated, truncated, info
+
+    def _is_crashed(self) -> bool:
+        pos = self.quad.get_position()
+        att = self.quad.get_attitude()
+        if pos[2] < 0.0:
+            return True
+        if abs(att[0]) > _TILT_LIMIT or abs(att[1]) > _TILT_LIMIT:
+            return True
+        if np.linalg.norm(pos) > _POS_LIMIT:
+            return True
+        return False
+
+    # ------------------------------------------------------------------
+    # Rendering
+    # ------------------------------------------------------------------
+
+    def render(self):
+        if self.render_mode == "viser":
+            _render_viser(self)
+
+    def close(self):
+        pass
+
+
+def _render_viser(env):
+    """Minimal viser rendering — push current pose to a viser frame."""
+    try:
+        import viser
+    except ImportError:  # pragma: no cover
+        return
+    server = getattr(env, "_viser_server", None)
+    handle = getattr(env, "_viser_handle", None)
+    if server is None:
+        server = viser.ViserServer(port=8082)
+        handle = server.scene.add_frame("/rl_quad")
+        env._viser_server = server
+        env._viser_handle = handle
+    pos = env.quad.get_position()
+    from viser import transforms as vtf
+
+    from drones_sim.math_utils import quat_to_rotation_matrix
+    wxyz = vtf.SO3.from_matrix(quat_to_rotation_matrix(env.quad.get_quaternion())).wxyz
+    handle.wxyz = tuple(wxyz)
+    handle.position = tuple(pos)

src/drones_sim/rl/observations.py

@@ -0,0 +1,40 @@
+"""Observation builders for the quadcopter RL environment.
+
+Each builder produces a ``dim``-dimensional numpy vector.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+class RelativeStateObs:
+    """Default observation: position error + velocity + quaternion + body rates.
+
+    17-D vector: [pos_err(3), vel(3), quat(4), omega(3), prev_action(4)].
+    All quantities in the world frame where applicable.
+    """
+
+    dim = 17
+
+    def __init__(self) -> None:
+        self._prev_action = np.zeros(4, dtype=np.float32)
+
+    def build(self, quad, task, action: NDArray | None = None) -> NDArray:
+        target = task.target_pos(quad)
+        pos_err = target - quad.get_position()
+        vel = quad.get_velocity()
+        quat = quad.get_quaternion()
+        omega = quad.get_angular_velocity()
+
+        if action is not None:
+            self._prev_action = action.astype(np.float32)
+
+        return np.concatenate([
+            pos_err.astype(np.float32),
+            vel.astype(np.float32),
+            quat.astype(np.float32),
+            omega.astype(np.float32),
+            self._prev_action,
+        ])

src/drones_sim/rl/reward.py

@@ -0,0 +1,69 @@
+"""Modular reward function for the quadcopter environment.
+
+Usage::
+
+    cfg = RewardConfig()
+    r = reward(quad, task, action, step_idx, cfg, prev_action)
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass
+class RewardConfig:
+    """Weights for the sum-of-terms reward function.
+
+    All terms are negative (cost) except survival and reach bonuses.
+    """
+
+    w_pos: float = 1.0       # position tracking (L2)
+    w_vel: float = 0.05      # velocity tracking
+    w_attitude: float = 0.1  # tilt penalty
+    w_action: float = 0.01   # control effort (L2)
+    w_action_d: float = 0.005  # control rate (Δu)
+    w_alive: float = 0.1     # survival bonus per step
+    w_reach: float = 10.0    # one-shot bonus for reaching target
+    w_crash: float = -10.0   # one-shot penalty for crashing
+    reach_radius: float = 0.1  # m — target-reached threshold
+
+
+def reward(quad, task, action: np.ndarray, step_idx: int,
+           cfg: RewardConfig, prev_action: np.ndarray | None = None) -> float:
+    """Compute per-step reward.
+
+    Parameters
+    ----------
+    quad: QuadcopterDynamics instance.
+    task: Task instance with ``target_pos(quad)`` and ``target_vel(quad)``.
+    action: The action that was just executed.
+    step_idx: Current step (0-based).
+    cfg: Reward weight configuration.
+    prev_action: Action from the previous step (for Δu penalty).
+    """
+    pos_err = float(np.linalg.norm(task.target_pos(quad) - quad.get_position()))
+    vel = quad.get_velocity()
+    att = quad.get_attitude()
+    tilt = float(np.sqrt(att[0]**2 + att[1]**2))
+
+    r_pos = -cfg.w_pos * pos_err
+    r_vel = -cfg.w_vel * float(np.linalg.norm(vel - task.target_vel(quad)))
+    r_att = -cfg.w_attitude * tilt
+    r_action = -cfg.w_action * float(np.sum(action**2))
+
+    r_action_d = 0.0
+    if prev_action is not None:
+        r_action_d = -cfg.w_action_d * float(np.linalg.norm(action - prev_action))
+
+    r_alive = cfg.w_alive
+
+    r = r_pos + r_vel + r_att + r_action + r_action_d + r_alive
+
+    # Terminal bonus — only at the last step
+    if pos_err < cfg.reach_radius:
+        r += cfg.w_reach
+
+    return float(r)