feat: Energy consumption

src/jdrones/data_models.py

@@ -263,6 +263,9 @@ class URDFModel(pydantic.BaseModel):
     g: float = 9.81
     """Acceleration due to gravity (m/s^2)"""
 
+    rho: float = 1.225
+    """Density of air at sea level (kg/m^3)"""
+
     l: float
     """Arm length (m)"""
 

src/jdrones/energy_model.py

@@ -0,0 +1,206 @@
+#  Copyright 2023 Jan-Hendrik Ewers
+#  SPDX-License-Identifier: GPL-3.0-only
+import abc
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+
+from jdrones.data_models import URDFModel
+from jdrones.envs.dronemodels import DronePlus
+
+
+class BaseEnergyModel(abc.ABC):
+    """
+    Base composite for energy and power implementations
+    """
+
+    dt: float
+
+    model: URDFModel
+    """Model parameters"""
+
+    def __init__(self, dt, model=DronePlus, *, _state=None):
+        self.dt = dt
+        self.model = model
+
+    @abc.abstractmethod
+    def power(self, state: dict[str, Any] | npt.ArrayLike):
+        """
+        Calculate the power usage of the system at the state
+
+        Parameters
+        ----------
+        state: Any
+
+        Returns
+        -------
+        list | float
+            The calculated power consumption
+        """
+        pass
+
+    def energy(self, state: dict[str, Any] | npt.ArrayLike):
+        """
+        :math:`E=P\\Delta t`
+
+        Parameters
+        ----------
+        state: Any
+            The state used to calculate the power
+
+        Returns
+        -------
+        list | float
+            The calculated energy
+        """
+
+        return self.power(state) * self.dt
+
+
+class StaticPropellerVariableVelocityEnergyModel(BaseEnergyModel):
+    """
+    Implement an energy model from [1,2,3] that is based on variable body velocity
+    magnitude. The assumption is that the UAV's propellers aren't changing in
+    RPM and is therefore an approximation as this clearly isn't the case in the
+    real world.
+
+    [1] D. Ebrahimi, S. Sharafeddine, P.-H. Ho, and C. Assi, ‘Autonomous UAV Trajectory
+        for Localizing Ground Objects: A Reinforcement Learning Approach’, IEEE
+        Transactions on Mobile Computing, vol. 20, no. 4, pp. 1312–1324, Apr. 2021,
+        doi: 10.1109/TMC.2020.2966989.
+    [2] A. Filippone, Flight performance of fixed and rotary wing aircraft, Ist ed. in
+        Elsevier aerospace engineering series. Amsterdam;
+        Boston: Butterworth-Heinemann, 2006.
+    [3] H. Sallouha, M. M. Azari, and S. Pollin, ‘Energy-Constrained UAV Trajectory
+        Design for Ground Node Localization’, in 2018 IEEE Global Communications
+        Conference (GLOBECOM), Dec. 2018, pp. 1–7. doi: 10.1109/GLOCOM.2018.8647530.
+    """
+
+    v_b: float
+    """Propeller RPM"""
+    K: float
+    """Constant related to drag induced by the propeller"""
+    F: float
+    """Constant related to the parasitic dragged"""
+    A: float
+    """Area of the drone"""
+
+    def __init__(
+        self,
+        dt: float,
+        model: URDFModel,
+        *,
+        v_b: float = 9000,
+        K: float = 1,
+        F: float = 0.1,
+        A: float = 1,
+    ):
+        super().__init__(dt, model)
+        self.v_b = v_b
+        self.K = K
+        self.F = F
+        self.A = A
+        self.m = self.model.mass
+        self.g = self.model.g
+        self.rho = self.model.rho
+
+    def p_blade(self, v: npt.ArrayLike):
+        """
+        Power required to turn the blades
+
+        .. math::
+            P_{blade} = K(1+3\\frac{v^2}{v_b^2})
+
+        Parameters
+        ----------
+        v : float | numpy.ndarray
+            Body velocity
+
+        Returns
+        -------
+        float | numpy.ndarray
+        """
+        return self.K * (1 + 3 * np.square(v) / np.square(self.v_b))
+
+    def p_parasite(self, v: npt.ArrayLike):
+        """
+        Power used to overcome the drag force
+
+        .. math::
+            P_{parasite} = \\frac{1}{2}\\rhov^3F
+
+        Parameters
+        ----------
+        v : float | numpy.ndarray
+            Body velocity
+
+        Returns
+        -------
+        float | numpy.ndarray
+        """
+        return 0.5 * self.rho * np.power(v, 3) * self.F
+
+    def v_i(self, v: npt.ArrayLike):
+        """
+        The mean propellers’ induced velocity in the forward flight
+
+        .. math::
+           v_i = \\sqrt{
+                \\frac{
+                -v^2 + \\sqrt{v^4+(\\frac{mg}{\rho A})^2}
+                }{2}
+                }
+
+        Parameters
+        ----------
+        v : float | numpy.ndarray
+            Body velocity
+
+        Returns
+        -------
+        float | numpy.ndarray
+        """
+        a = -np.square(v)
+        b = np.power(v, 4) + np.square(self.m * self.g / (self.rho * self.A))
+
+        return np.sqrt((a + np.sqrt(b)) / 2)
+
+    def p_induced(self, v: npt.ArrayLike):
+        """
+        Power required to lift the UAV and overcome the drag caused by gravity
+
+        .. math::
+            P_{parasite} = \\frac{1}{2}\\rhov^3F
+
+        Parameters
+        ----------
+        v : float | numpy.ndarray
+            Body velocity
+
+        Returns
+        -------
+        float | numpy.ndarray
+        """
+
+        return self.m * self.g * self.v_i(v)
+
+    def power(self, v: npt.ArrayLike):
+        """
+        Total power required at the system velocity
+        
+        .. math::
+        
+            P_{total} = P_{blade} + P_{parasite} + P_{induced}
+
+
+        Parameters
+        ----------
+        v: float | numpy.ndarray
+
+        Returns
+        -------
+        list | float
+            The calculated total power consumption
+        """
+        return self.p_blade(v) + self.p_parasite(v) + self.p_induced(v)

src/jdrones/wrappers.py

@@ -0,0 +1,51 @@
+from typing import SupportsFloat, Any
+
+import gymnasium
+import numpy as np
+from gymnasium.core import WrapperActType, WrapperObsType
+
+from jdrones.envs.base.basedronenev import BaseDroneEnv
+from jdrones.energy_model import (
+    StaticPropellerVariableVelocityEnergyModel,
+    BaseEnergyModel,
+)
+
+
+class EnergyCalculationWrapper(gymnasium.Wrapper):
+    """
+    Wrap a drone env to get energy consumption calculations in the returned `info`.
+    """
+
+    def __init__(
+        self,
+        env: BaseDroneEnv,
+        energy_model: type[
+            BaseEnergyModel
+        ] = StaticPropellerVariableVelocityEnergyModel,
+    ):
+        super().__init__(env)
+        if hasattr(self.env, "model"):
+            model = self.env.model
+        elif hasattr(self.env, "env"):
+            model = self.env.env.model
+        else:
+            raise ValueError("Could not find model information within the wrapped "
+                             "environment")
+        self.energy_calculation = energy_model(env.dt, model)
+
+    def step(
+        self, action: WrapperActType
+    ) -> tuple[WrapperObsType, SupportsFloat, bool, bool, dict[str, Any]]:
+        obs, reward, term, trunc, info = super().step(action)
+
+        if hasattr(self.env, "state"):
+            vel = self.env.state.vel
+        elif hasattr(self.env, "env"):
+            vel = self.env.env.state.vel
+        else:
+            raise ValueError("Could not find velocity information within the wrapped "
+                             "environment")
+        speed = np.linalg.norm(vel)
+        info["energy"] = self.energy_calculation.energy(speed)
+
+        return obs, reward, term, trunc, info

tests/conftest.py

@@ -83,6 +83,11 @@ def tau_Q(request):
 def drag_coeffs(request):
     return request.param
 
+@pytest.fixture(params=[1.225])
+def rho(request):
+    return request.param
+
+
 
 @pytest.fixture(params=["droneplus.urdf"])
 def filepath(request):
@@ -162,6 +167,7 @@ def urdfmodel(
     tau_Q,
     drag_coeffs,
     mass,
+        rho,
     filepath,
     mixing_matrix,
 ):
@@ -178,6 +184,7 @@ def urdfmodel(
         filepath=str(filepath),
         mixing_matrix=mixing_matrix,
         max_vel_ms=1,
+        rho=rho
     )
 
 

tests/envs/position/test_stress.py

@@ -2,6 +2,8 @@
 #  SPDX-License-Identifier: GPL-3.0-only
 import pytest
 
+from jdrones.wrappers import EnergyCalculationWrapper
+
 
 def run(T, dt, env):
     trunc = False
@@ -30,7 +32,6 @@ def run(T, dt, env):
 )
 T_PARAM = pytest.mark.parametrize("T", [2, 10, 100])
 
-
 @pytest.mark.slow_integration
 @AS_PARAM
 @T_PARAM
@@ -50,3 +51,11 @@ def test_long_poly_position(T, dt, fifthorderpolypositiondroneenv):
 @T_PARAM
 def test_long_lookahead_position(T, dt, fifthorderpolypositionlookaheaddroneenv):
     run(T, dt, fifthorderpolypositionlookaheaddroneenv)
+
+@pytest.mark.slow_integration
+@AS_PARAM
+@pytest.mark.parametrize("T", [2])
+@pytest.mark.parametrize("wrapper",[EnergyCalculationWrapper])
+def test_wrappers(T, dt, wrapper, firstorderploypositiondroneenv):
+    env = wrapper(firstorderploypositiondroneenv)
+    run(T, dt, env)