Added motor simulation, gym env interfaces;

env.py: Dynamics- and Speed-based environment classes,
vehicle.py: Motor parameters,
simulation.py: Motor simulation,
helpers.py: Added function to learn speed-to-voltage conversion given a motor

env.py

+"""
+This module defines OpenAI Gym compatible classes based on the Multirotor class.
+"""
+
+import numpy as np
+import gym
+
+from .simulation import Multirotor
+
+
+class BaseMultirotorEnv(gym.Env):
+
+
+    def __init__(self, vehicle: Multirotor=None) -> None:
+        # pos, vel, att, ang vel
+        self.observation_space = gym.spaces.Box(
+            low=-np.inf,
+            high=np.inf,
+            dtype=np.float32,
+            shape=(12,)
+        )
+        self.vehicle = vehicle
+        self.state : np.ndarray = None
+
+
+    def reset(self):
+        if self.vehicle is not None:
+            self.state = self.vehicle.reset()
+        return self.state
+
+
+    def reward(self, state, action, nstate):
+        raise NotImplementedError
+
+
+    def step(self, action: np.ndarray) -> tuple[np.ndarray, float, bool, dict]:
+        raise NotImplementedError
+
+
+
+class DynamicsMultirotorEnv(BaseMultirotorEnv):
+
+
+    def __init__(self, vehicle: Multirotor=None, allocate: bool=False) -> None:
+        super().__init__(vehicle=vehicle)
+        
+        self.action_space = gym.spaces.Box(
+            low=-np.inf,
+            high=np.inf,
+            dtype=np.float32,
+            shape=(6,)  # 3 forces, 3 torques
+        )
+        self.allocate = allocate
+
+
+    def step(self, action: np.ndarray):
+        if self.allocate:
+            speeds = self.vehicle.allocate_control(action[0], action[3:6])
+            forces, torques = self.vehicle.get_forces_torques(speeds, self.vehicle.state)
+            action = np.concatenate(forces, torques)
+        state = self.vehicle.step_dynamics(u=action)
+        return state, None, None, None
+
+
+
+class SpeedsMultirotorEnv(gym.Env):
+
+
+    def __init__(self, vehicle: Multirotor=None) -> None:
+        super().__init__(vehicle=vehicle)
+        
+        self.action_space = gym.spaces.Box(
+            low=-np.inf,
+            high=np.inf,
+            dtype=np.float32,
+            shape=(len(vehicle.propellers),)  # action for each propeller
+        )
+
+
+    def step(self, action: np.ndarray):
+        state = self.vehicle.step_speeds(u=action)
+        return state, None, None, None
\ No newline at end of file

helpers.py

@@ -81,6 +81,17 @@ def learn_thrust_coefficient(
 
 
 
+def learn_speed_voltage_scaling(
+    speed_fn: Callable[[float], float], domain: Tuple=(0,20)
+) -> float:
+    signals = np.linspace(domain[0], domain[1], num=10)
+    speeds = np.zeros_like(signals)
+    for i, signal in enumerate(signals):
+        speeds[i] = speed_fn(signal)
+    return np.polyfit(signals, speeds, 1)[0]
+
+
+
 def moment_of_inertia_disk(m: float, r: float) -> float:
     return 0.5 * m * r**2
 
@@ -94,6 +105,6 @@ def control_allocation_matrix(params: VehicleParams) -> Tuple[np.ndarray, np.nda
         alloc[0, i] = p.k_thrust                       # vertical force
         alloc[1, i] = p.k_thrust * y[i]                 # torque about x-axis
         alloc[2, i] = p.k_thrust * (-x[i])              # torque about y-axis
-        alloc[3, i] = p.k_torque * params.clockwise[i]    # torque about z-axis
+        alloc[3, i] = p.k_drag * params.clockwise[i]    # torque about z-axis
     alloc_inverse = np.linalg.pinv(alloc)
     return alloc, alloc_inverse
\ No newline at end of file

simulation.py

-from typing import List, Tuple
+from typing import Callable, List, Tuple
 from copy import deepcopy
 
 import numpy as np
 from numpy import (cos, sin)
-from scipy.integrate import odeint
+from scipy.integrate import odeint, trapezoid
 
-from .vehicle import MotorParams, PropellerParams, SimulationParams, VehicleParams
+from .vehicle import MotorParams, PropellerParams, SimulationParams, VehicleParams, BatteryParams
 from .coords import body_to_inertial, direction_cosine_matrix, rotating_frame_derivative, angular_to_euler_rate
 from .physics import thrust, torque, apply_forces_torques
 from .helpers import control_allocation_matrix
@@ -29,27 +29,30 @@ class Propeller:
         self.use_thrust_constant = use_thrust_constant
 
         if use_thrust_constant:
-            self.thrust = self._thrust_constant
+            self.thrust: Callable = self._thrust_constant
         else:
-            self.thrust = self._thrust_physics
+            self.thrust: Callable = self._thrust_physics
+        
+        if params.motor is not None:
+            self.motor = Motor(params.motor, self.simulation)
+        else:
+            self.motor: Motor = None
 
 
     def reset(self):
         self.speed = 0.
+        if self.motor is not None:
+            self.motor.reset()
 
 
-    def step(self, u: float) -> float:
-        return self.apply_speed(u)
-
-
-    def apply_speed(self, speed: float) -> float:
+    def apply_speed(self, u: float) -> float:
         """
         Calculate the actual speed of the propeller after the speed signal is
         given. This method is *pure* and does not change the state of the propeller.
 
         Parameters
         ----------
-        speed : float
+        u : float
             Radians per second speed command
 
         Returns
@@ -57,17 +60,19 @@ class Propeller:
         float
             The actual speed
         """
-        return speed
+        if self.motor is not None:
+            return self.motor.apply_speed(u)
+        return u
 
 
-    def step(self, speed: float) -> float:
+    def step(self, u: float) -> float:
         """
         Step through the speed command. This method changes the state of the 
         propeller.
 
         Parameters
         ----------
-        speed : float
+        u : float
             Speed command in radians per second.
 
         Returns
@@ -75,7 +80,11 @@ class Propeller:
         float
             The actual speed achieved.
         """
-        self.speed = speed
+        if self.motor is not None:
+            self.speed = self.motor.step(u)
+        else:
+            self.speed = u
+        return self.speed
 
 
     def _thrust_constant(self, speed,airstream_velocity: np.ndarray=np.zeros(3)) -> float:
@@ -112,13 +121,56 @@ class Motor:
         self.simulation = simulation
         self.speed: float = 0.
         "Radians per second"
+        self._net_torques = np.zeros(2)
 
 
     def reset(self) -> float:
         self.speed = 0.
+        self.voltage = 0.
+        self.current = 0.
+        self._net_torques *= 0
+
+
+    def apply_speed(self, u: float) -> float:
+        voltage = self.params.speed_voltage_scaling * u
+        current = (voltage - self.speed * self.params.k_emf) / self.params.resistance
+        torque = self.params.k_torque * current
+        # Subtract drag torque and dynamic friction from electrical torque
+        net_torque = torque - \
+                     self.params.k_df * self.speed - \
+                     self.params.k_drag * self.speed**2
+        self._net_torques[0] = self._net_torques[1]
+        self._net_torques[1] = net_torque
+        return self.speed + \
+        trapezoid(
+            self._net_torques / self.params.moment_of_inertia,
+            dx=self.simulation.dt
+        )
 
 
     def step(self, u: float) -> float:
+        self.voltage = self.params.speed_voltage_scaling * u
+        self.current = (self.voltage - self.speed * self.params.k_emf) / self.params.resistance
+        self.speed = self.apply_speed(u)
+        return self.speed
+
+
+
+class Battery:
+    """
+    Models the state of charge of the battery of the Multirotor.
+    """
+
+    def __init__(self, params: BatteryParams, simulation: SimulationParams) -> None:
+        self.params = deepcopy(params)
+        self.simulation = simulation
+
+
+    def reset(self):
+        pass
+
+
+    def step(self):
         pass
 
 
@@ -220,7 +272,7 @@ class Multirotor:
             torque_vec[:, i] = torque(
                 self.propeller_vectors[:,i], thrust_vec[:,i],
                 prop.params.moment_of_inertia, angular_acc,
-                prop.params.k_torque, speed,
+                prop.params.k_drag, speed,
                 clockwise
             )
         forces = thrust_vec.sum(axis=1)
@@ -254,20 +306,17 @@ class Multirotor:
 
     def step_dynamics(self, u: np.ndarray):
         self.t += self.simulation.dt
-        # TODO: Explore RK45() or solve_ivp() functions from scipy.integrate?
         self.state = odeint(
             self.dxdt_dynamics, self.state, (0, self.simulation.dt), args=(u,),
             rtol=1e-4, atol=1e-4
         )[-1]
         self.state = np.around(self.state, 4)
-        for u_, prop in zip(u, self.propellers):
-            prop.step(u_)
+        # TODO: inverse solve for speed = forces to set propeller speeds
         return self.state
 
 
     def step_speeds(self, u: np.ndarray):
         self.t += self.simulation.dt
-        # TODO: Explore RK45() or solve_ivp() functions from scipy.integrate?
         self.state = odeint(
             self.dxdt_speeds, self.state, (0, self.simulation.dt), args=(u,),
             rtol=1e-4, atol=1e-4

vehicle.py

@@ -8,22 +8,32 @@ import numpy as np
 @dataclass
 class MotorParams:
 
-    r: float
-    "Motor resistance"
-    moment_of_inertia: float
+    resistance: float
+    "Motor resistance r. Current = (voltage - back_emf) / resistance"
+     # See: http://learningrc.com/motor-kv/, http://web.mit.edu/first/scooter/motormath.pdf
+    k_motor: float = None
+    "Motor constant k_m, where speed = k_m * back_emf, and k_m = 1/k_e"
+    k_emf: float = None
+    "Back-EMF constant k_e, relates motor speed induced voltage, back_emf = k_e * omega"
+    k_torque: float = None
+    "Torque constant k_q, where torque Q = k_q * current. Equal to k_e"
+    k_drag: float = None
+    "Aerodynamic drag coefficient, where torque = k_drag * omega^2"
+    moment_of_inertia: float = 0.
     "Moment of inertia about rotational axis"
-    d_f: float
-    "Viscous damping coefficient"
-    static_friction: float
-    
+    k_df: float = 0.
+    "Viscous damping coefficient. Torque = d_f * speed"
+    static_friction: float = 0.
+    speed_voltage_scaling: float = 1.
+    """Scaling constant to convert speed signal (rad/s) into speed controller voltage (V).
+    If 1, means input action is same as Voltage"""
 
-    # See: http://learningrc.com/motor-kv/, http://web.mit.edu/first/scooter/motormath.pdf
-    k_m: float
-    "Motor constant, where omega = k_m * back_emf, and k_m = 1/k_e"
-    k_e: float = None
-    "Back-EMF constant, relates motor speed induced voltage, back_emf = k_e * omega"
-    k_q: float = None
-    "Torque constant k_q, where torque Q = k_q * current. Equal to k_e"
+
+    def __post_init__(self):
+        # See: https://www.motioncontroltips.com/faq-difference-between-torque-back-emf-motor-constant/
+        # For an ideal square-wave motor, torque and back-emf constants are same
+        self.k_torque = self.k_torque or self.k_emf
+        self.k_emf = self.k_emf or self.k_torque
 
 
 @dataclass
@@ -48,8 +58,8 @@ class PropellerParams:
 
     k_thrust: float = None
     "Propeller's aerodynamic thrust coefficient, where thrust =  k_thrust * angular velocity^2"
-    k_torque: float = None
-    "Torque constant, where torque = k_torque * angular velocity^2"
+    k_drag: float = None
+    "Torque constant or drag coefficient, where torque = k_drag * angular velocity^2"
 
     motor: MotorParams = None
     "The parameters of the motor to simulate speed, otherwise instantaneous."
@@ -66,6 +76,8 @@ class PropellerParams:
         "Pitch angle at root of blade"
         self.theta1 = -4 / 3 * np.arctan2(self.pitch, 2 * np.pi * 3/4 * self.diameter/2)
         "Change in pitch angle towards tip of blade"
+        if self.motor is not None:
+            self.motor.k_drag = self.k_drag
         # Pitch angel is reduced to allow for even lift as blade velocity increases
         # with increasing radius. Assuming linear reduction from root to tip.
 
@@ -106,3 +118,11 @@ class SimulationParams:
     """Gravitational acceleration"""
     rho: float = 1.225
     "Air density kg/m^3 at MSL"
+
+
+
+@dataclass
+class BatteryParams:
+
+    max_voltage: float = 20
+    "Maximum voltage of the battery"