drones-sim 0.2.0__py3-none-any.whl

drones_sim/rl/env.py

@@ -0,0 +1,195 @@
+"""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)
+        # Start at the task target position so the drone doesn't begin
+        # on the ground where a single negative-thrust step is fatal.
+        # Small random offset (±0.5 m) forces the policy to generalize.
+        target_pos = self.task.target_pos(self.quad)
+        offset = self._rng.uniform(-0.5, 0.5, size=3)
+        self.quad.reset(position=target_pos + offset)
+        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.
+        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):
+        # Let LQRResidualAction know the task target
+        self.action_param._task_target = self.task.target_pos(self.quad)
+        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)

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,
+        ])

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)

drones_sim/rl/tasks.py

@@ -0,0 +1,73 @@
+"""Task definitions for the quadcopter RL environment."""
+
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+class HoverTask:
+    """Hover at a fixed target position — the simplest baseline task."""
+
+    def __init__(self, target: tuple[float, float, float] = (0.0, 0.0, 2.0)) -> None:
+        self._target = np.array(target, dtype=float)
+
+    def reset(self, quad, rng: np.random.Generator | None = None) -> None:
+        pass
+
+    def target_pos(self, _quad) -> NDArray:
+        return self._target
+
+    def target_vel(self, _quad) -> NDArray:
+        return np.zeros(3)
+
+
+class WaypointTask:
+    """Sequence of waypoints — advance when within reach_radius."""
+
+    def __init__(
+        self,
+        waypoints: list[tuple[float, float, float]],
+        reach_radius: float = 0.2,
+    ) -> None:
+        self.waypoints = [np.array(w, dtype=float) for w in waypoints]
+        self.reach_radius = reach_radius
+        self._idx = 0
+
+    def reset(self, quad, rng: np.random.Generator | None = None) -> None:
+        self._idx = 0
+
+    def target_pos(self, quad) -> NDArray:
+        if self._idx < len(self.waypoints):
+            dist = np.linalg.norm(self.waypoints[self._idx] - quad.get_position())
+            if dist < self.reach_radius and self._idx < len(self.waypoints) - 1:
+                self._idx += 1
+            return self.waypoints[self._idx]
+        return self.waypoints[-1]
+
+    def target_vel(self, _quad) -> NDArray:
+        return np.zeros(3)
+
+
+class TrackingTask:
+    """Track a pre-generated TrajectoryData (circular / minimum-snap).
+
+    The closest-in-time reference position is returned at each step.
+    """
+
+    def __init__(self, traj) -> None:
+        self.traj = traj
+        self._step = 0
+
+    def reset(self, quad, rng: np.random.Generator | None = None) -> None:
+        self._step = 0
+
+    def target_pos(self, _quad, advance: bool = True) -> NDArray:
+        if advance:
+            self._step += 1
+        idx = min(self._step, len(self.traj.position) - 1)
+        return self.traj.position[idx]
+
+    def target_vel(self, _quad) -> NDArray:
+        idx = min(self._step, len(self.traj.position) - 1)
+        return self.traj.velocity[idx] if hasattr(self.traj, 'velocity') else np.zeros(3)

drones_sim/sensors/__init__.py

@@ -0,0 +1,3 @@
+from .gps import GPSConfig, GPSData, GPSSimulator  # noqa: F401
+from .imu import IMUConfig, IMUData, IMUMeasurement, IMUSimulator  # noqa: F401
+from .models import SensorNoiseModel, TemperatureModel  # noqa: F401

drones_sim/sensors/gps.py

@@ -0,0 +1,158 @@
+"""GPS sensor simulator following the IMUSimulator pattern.
+
+Simulates a low-rate GNSS receiver producing noisy 3-D position and
+optionally 3-D velocity measurements from a ground-truth trajectory.
+
+Typical usage
+-------------
+>>> gps = GPSSimulator()
+>>> gps_data = gps.simulate(traj)
+>>> for k in range(len(gps_data.t)):
+...     if k % gps_update_steps == 0:
+...         ekf.correct_position(gps_data.position[k])
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..trajectory import TrajectoryData
+
+
+@dataclass
+class GPSConfig:
+    """Configuration parameters for the GPS sensor model."""
+    position_noise_std: float = 1.0      # m  (1-sigma, each axis)
+    velocity_noise_std: float = 0.1      # m/s (1-sigma, each axis)
+    update_rate: float = 5.0             # Hz — GPS fix rate
+    dropout_probability: float = 0.0     # probability of fix dropout per epoch
+
+
+@dataclass
+class GPSData:
+    """Container for GPS measurement sequences.
+
+    Attributes
+    ----------
+    t:
+        Time stamps of GPS epochs [s], shape (K,).
+    position:
+        Noisy 3-D position [m], shape (K, 3).
+    velocity:
+        Noisy 3-D velocity [m/s], shape (K, 3).
+    valid:
+        Boolean mask: True when the fix is valid (no dropout), shape (K,).
+    """
+    t: NDArray
+    position: NDArray   # (K, 3)
+    velocity: NDArray   # (K, 3)
+    valid: NDArray      # (K,)  bool
+
+
+class GPSSimulator:
+    """Simulate a low-rate GPS receiver from a ground-truth trajectory.
+
+    The noise model adds independent Gaussian noise to each position and
+    velocity axis at the configured update rate.  An optional dropout
+    probability marks epochs as invalid (simulating urban-canyon outages).
+
+    Parameters
+    ----------
+    config:
+        Sensor configuration.  Defaults to ``GPSConfig()``.
+    seed:
+        Random seed for reproducibility.  ``None`` for non-deterministic.
+    """
+
+    def __init__(
+        self,
+        config: GPSConfig | None = None,
+        seed: int | None = None,
+    ) -> None:
+        self.config = config or GPSConfig()
+        self._seed = seed
+        self._rng = np.random.default_rng(seed)
+
+    def reset(self, seed: int | None = None) -> None:
+        """Reset the receiver's local random stream for a new episode."""
+        self._rng = np.random.default_rng(self._seed if seed is None else seed)
+
+    def simulate(self, traj: TrajectoryData) -> GPSData:
+        """Generate GPS measurements from a full ground-truth trajectory.
+
+        Epochs are spaced ``1 / update_rate`` seconds apart, aligned with
+        the nearest sample in ``traj.t``.
+
+        Parameters
+        ----------
+        traj:
+            Ground-truth trajectory.
+
+        Returns
+        -------
+        GPSData
+            GPS measurement sequences at the GPS update rate.
+        """
+        cfg = self.config
+        dt_gps = 1.0 / cfg.update_rate
+        t_max  = float(traj.t[-1])
+        t_epochs = np.arange(0.0, t_max + dt_gps * 0.5, dt_gps)
+
+        K = len(t_epochs)
+        pos_meas = np.zeros((K, 3))
+        vel_meas = np.zeros((K, 3))
+        valid    = np.ones(K, dtype=bool)
+
+        for i, te in enumerate(t_epochs):
+            # Nearest sample index
+            idx = int(np.argmin(np.abs(traj.t - te)))
+
+            # Add Gaussian noise
+            pos_noise = self._rng.normal(0.0, cfg.position_noise_std, 3)
+            vel_noise = self._rng.normal(0.0, cfg.velocity_noise_std, 3)
+
+            pos_meas[i] = traj.position[idx] + pos_noise
+            vel_meas[i] = traj.velocity[idx] + vel_noise
+
+            # Optional dropout
+            if cfg.dropout_probability > 0.0:
+                valid[i] = self._rng.random() >= cfg.dropout_probability
+
+        return GPSData(
+            t=t_epochs,
+            position=pos_meas,
+            velocity=vel_meas,
+            valid=valid,
+        )
+
+    def step(
+        self,
+        true_position: NDArray,
+        true_velocity: NDArray,
+    ) -> tuple[NDArray, NDArray, bool]:
+        """Single-step GPS measurement for closed-loop simulation.
+
+        Parameters
+        ----------
+        true_position:
+            Ground-truth 3-D position [m].
+        true_velocity:
+            Ground-truth 3-D velocity [m/s].
+
+        Returns
+        -------
+        pos_meas, vel_meas, is_valid
+        """
+        cfg = self.config
+        pos_noise = self._rng.normal(0.0, cfg.position_noise_std, 3)
+        vel_noise = self._rng.normal(0.0, cfg.velocity_noise_std, 3)
+        is_valid  = (cfg.dropout_probability == 0.0 or
+                     self._rng.random() >= cfg.dropout_probability)
+        return (
+            true_position + pos_noise,
+            true_velocity + vel_noise,
+            is_valid,
+        )