Merge pull request #44 from iwishiwasaneagle/dev-improve-fifth-order

.github/workflows/CI.yml

@@ -62,10 +62,11 @@ jobs:
     - name: Install dependencies
       if: steps.cache-virtualenv.outputs.cache-hit != 'true'
       run: |
-        pip install . -r requirements.txt -r tests/requirements.txt
+        pip install -r requirements.txt -r tests/requirements.txt
 
     - name: Run tests
       run: |
+        PYTHONPATH=src:tests \
         pytest tests \
         --cov-report=xml \
         --cov-branch \

README.md

@@ -23,7 +23,8 @@ The environment documentation can be found [here](https://jdrones.janhendrikewer
 
 ### Position
 1. `FirstOrderPolyPositionDroneEnv-v0` [:link:](https://jdrones.janhendrikewers.uk/envs.html#firstorderpolypositiondroneenv)
-2`FifthOrderPolyPositionDroneEnv-v0` [:link:](https://jdrones.janhendrikewers.uk/envs.html#fifthorderpolypositiondroneenv)
+2. `FifthOrderPolyPositionDroneEnv-v0` [:link:](https://jdrones.janhendrikewers.uk/envs.html#fifthorderpolypositiondroneenv)
+3. `FifthOrderPolyPositionWithLookAheadDroneEnv-v0` [:link:](https://jdrones.janhendrikewers.uk/envs.html#fifthorderpolypositionwithlookaheaddroneenv)
 
 ## Development
 

docs/envs.rst

@@ -73,3 +73,10 @@ FifthOrderPolyPositionDroneEnv
    :members:
    :undoc-members:
    :show-inheritance:
+
+FifthOrderPolyPositionDroneEnvWithLookAheadDroneEnv
+---------------------------------------------------
+.. autoclass:: jdrones.envs.FifthOrderPolyPositionDroneEnvWithLookAheadDroneEnv
+   :members:
+   :undoc-members:
+   :show-inheritance:

src/jdrones/__init__.py

@@ -24,6 +24,10 @@
     "FifthOrderPolyPositionDroneEnv-v0",
     entry_point="jdrones.envs:FifthOrderPolyPositionDroneEnv",
 )
+register(
+    "FifthOrderPolyPositionWithLookAheadDroneEnv-v0",
+    entry_point="jdrones.envs:FifthOrderPolyPositionWithLookAheadDroneEnv",
+)
 
 
 __version__ = "unknown"

src/jdrones/envs/__init__.py

@@ -7,6 +7,7 @@
 from .dronemodels import DronePlus
 from .lqr import LQRDroneEnv
 from .position import FifthOrderPolyPositionDroneEnv
+from .position import FifthOrderPolyPositionWithLookAheadDroneEnv
 from .position import FirstOrderPolyPositionDroneEnv
 
 __all__ = [
@@ -17,5 +18,6 @@
     "FirstOrderPolyPositionDroneEnv",
     "DronePlus",
     "FifthOrderPolyPositionDroneEnv",
+    "FifthOrderPolyPositionWithLookAheadDroneEnv",
     "BaseControlledEnv",
 ]

src/jdrones/envs/position.py

@@ -18,6 +18,8 @@
 from jdrones.trajectory import FifthOrderPolynomialTrajectory
 from jdrones.trajectory import FirstOrderPolynomialTrajectory
 from jdrones.types import PositionAction
+from jdrones.types import PositionVelocityAction
+from jdrones.types import VEC3
 
 
 class BasePositionDroneEnv(gymnasium.Env, abc.ABC):
@@ -49,43 +51,6 @@ def __init__(
             low=np.array([0, 0, 1]), high=np.array([10, 10, 10])
         )
 
-    def reset(
-        self,
-        *,
-        seed: int | None = None,
-        options: dict[str, Any] | None = None,
-    ) -> tuple[States, dict[str, Any]]:
-        super().reset(seed=seed, options=options)
-
-        obs, _ = self.env.reset(seed=seed, options=options)
-
-        return States([np.copy(obs)]), {}
-
-    @abc.abstractmethod
-    def step(
-        self, action: PositionAction
-    ) -> tuple[States, float, bool, bool, dict[str, Any]]:
-        """
-        A step from the viewpoint of a
-        :class:`~jdrones.envs.position.BasePositionDroneEnv` is making the drone
-        fly from its current position :math:`A` to the target position :math:`B`.
-
-
-        Parameters
-        ----------
-        action : float,float,float
-            Target coordinates :math:`(x,y,z)`
-
-        Returns
-        -------
-            states: jdrones.data_models.States
-            reward: float
-            term: bool
-            trunc: bool
-            info: dict
-        """
-        pass
-
     @staticmethod
     def get_reward(states: States) -> float:
         """
@@ -148,11 +113,78 @@ def update_u_from_traj(u: State, traj: BasePolynomialTrajectory, t: float):
             u.vel = traj.velocity(t)
         return u
 
+    @staticmethod
+    def _validate_action_input(
+        action: PositionAction | PositionVelocityAction,
+    ) -> State:
+        """
+        Validate the action input. Has to either be only position, or position and
+        velocity
+
+        Parameters
+        ----------
+        action : VEC3 | tuple[VEC3, VEC3]
+            Target coordinates :math:`(x,y,z)` or target coordinates and velocity (
+            :math:`(v_x,v_y,v_z)`)
+
+        Returns
+        -------
+        State
+            The valid and converted action as a :class:`~jdrones.data_models.State`
+
+        """
+        action_as_state = State()
+        action_np = np.asarray(action).squeeze()
+        if action_np.shape == (3,):
+            action_as_state.pos = action_np
+        elif action_np.shape == (2, 3):
+            action_as_state.pos = action_np[0]
+            action_as_state.vel = action_np[1]
+        elif action_np.shape == (6,):
+            action_as_state.pos = action_np[:3]
+            action_as_state.vel = action_np[3:]
+        else:
+            raise ValueError(
+                f"Incorrect shape {action_np.shape}. Expected either " f"(3,) or (2,3)"
+            )
+        return action_as_state
+
+    def reset(
+        self,
+        *,
+        seed: int | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> tuple[States, dict[str, Any]]:
+        super().reset(seed=seed, options=options)
+
+        obs, _ = self.env.reset(seed=seed, options=options)
+
+        return States([np.copy(obs)]), {}
+
     def step(
-        self, action: PositionAction
+        self, action: PositionAction | PositionVelocityAction
     ) -> tuple[States, float, bool, bool, dict[str, Any]]:
-        action_as_state = State()
-        action_as_state.pos = action
+        """
+        A step from the viewpoint of a
+        :class:`~jdrones.envs.position.BasePositionDroneEnv` is making the drone
+        fly from its current position :math:`A` to the target position :math:`B`.
+
+
+        Parameters
+        ----------
+        action : VEC3 | tuple[VEC3, VEC3]
+            Target coordinates :math:`(x,y,z)` or target coordinates and velocity (
+            :math:`(v_x,v_y,v_z)`)
+
+        Returns
+        -------
+            states: jdrones.data_models.States
+            reward: float
+            term: bool
+            trunc: bool
+            info: dict
+        """
+        action_as_state = self._validate_action_input(action)
 
         term, trunc, info = False, False, {}
         if np.allclose(action, self.env.state.pos):
@@ -173,7 +205,7 @@ def step(
         while not (term or trunc):
             t = next(counter) * self.dt
             if t > traj.T:
-                u.pos = action
+                u.pos = action_as_state.pos
                 u.vel = (0, 0, 0)
             else:
                 u = self.update_u_from_traj(u, traj, t)
@@ -303,3 +335,136 @@ def calc_traj(
             T=T,
         )
         return t
+
+
+class FifthOrderPolyPositionWithLookAheadDroneEnv(BasePositionDroneEnv):
+    """
+    Supplements the :class:`jdrones.env.FifthOrderPolyPositionDroneEnv` by including
+    the next waypoints position in the trajectory generation calculation.
+
+    >>> import jdrones
+    >>> import gymnasium
+    >>> gymnasium.make("FifthOrderPolyPositionWithLookAheadDroneEnv-v0")
+    <OrderEnforcing<PassiveEnvChecker<FifthOrderPolyPositionWithLookAheadDroneEnv<FifthOrderPolyPositionWithLookAheadDroneEnv-v0>>>>
+    """
+
+    def __init__(
+        self,
+        model: URDFModel = DronePlus,
+        initial_state: State = None,
+        dt: float = 1 / 240,
+        env: LQRDroneEnv = None,
+    ):
+        if env is None:
+            env = FifthOrderPolyPositionDroneEnv(
+                model=model, initial_state=initial_state, dt=dt, env=env
+            )
+        super().__init__(model=model, initial_state=initial_state, dt=dt, env=env)
+
+        self.action_space = spaces.Box(
+            low=np.array([[0, 0, 1], [0, 0, 1]]), high=10, shape=(2, 3)
+        )
+
+    def reset(
+        self,
+        *,
+        seed: int | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> tuple[States, dict[str, Any]]:
+        super().reset(seed=seed, options=options)
+
+        obs, _ = self.env.reset(seed=seed, options=options)
+
+        return obs, {}
+
+    @staticmethod
+    def calc_v_at_B(A: VEC3, B: VEC3, C: VEC3, *, V: float, N: float = 3):
+        """
+        For the velocity vector at :math:`B`, it is ideal to create a smooth
+        transition between the segments. One solution is to reduce velocity to 0.
+        However, this is very inefficient and not representative of real-world flight,
+        especially for the case where there are multiple waypoints in a straight line.
+        The preferred, albeit slightly more computationally expensive, solution is to
+        use the waypoints before and after to calculate a smooth velocity at :math:`B`.
+        This is calculated by
+
+        .. math::
+            \\begin{gather}
+                \\theta =
+                    \\frac{(-\\vec r_{AB}) \\cdot \\vec r_{BC}}
+                    {|\\vec r_{AB}| |\\vec r_{BC}|}, \\\\
+                \\chi = \\frac{(1-\\theta)}{2}, \\\\
+                \\vec v_{B,-} = \\vec v_{B,+} = V \\cdot
+                    \\frac{\\vec r_{AC}}
+                    {| \\vec r_{AC} |}
+                    \\cdot \\chi^N,
+            \\end{gather}
+
+
+        Parameters
+        ----------
+        A : VEC3
+            Current position
+        B : VEC3
+            Target position
+        C : VEC3
+            Position after going through :math:`B`
+        V : float
+            Scaling factor for the magnitude of :math:`\\vec v_B`
+        N : float
+            Scaling factor for the magnitude of :math:`\\chi`.
+            :math:`N = 3` for an aggressive reduction in velocity for non-straight line
+            waypoints to closely follow the path. :math:`N = 1` would be more suitable
+            for cases where following the path is not as important.
+            (Default = 3)
+
+        Returns
+        -------
+        VEC3
+            The velocity at :math:`B`
+        """
+
+        rAC = C - A
+        rBA = A - B
+        rBC = C - B
+        rAC_unit = rAC / np.clip(np.linalg.norm(rAC), 1e-6, np.inf)
+        theta = rBA.dot(rBC) / (np.linalg.norm(rBA) * np.linalg.norm(rBC))
+        chi = np.power((1 - theta) / 2, N)
+        return V * rAC_unit * chi
+
+    def step(
+        self, action: tuple[PositionAction, PositionAction]
+    ) -> tuple[States, float, bool, bool, dict[str, Any]]:
+        """
+        A step from the viewpoint of a
+        :class:`~jdrones.envs.position.
+        FifthOrderPolyPositionWithLookAheadDroneEnv` is making the drone
+        fly from its current position :math:`A` to the target position :math:`C` via
+        :math:`B`.
+
+        .. seealso::
+            :meth:`~jdrones.envs.position.
+            FifthOrderPolyPositionDroneEnvWithLookAheadDroneEnv.calc_v_at_B`
+
+        Parameters
+        ----------
+        action : tuple[VEC3, VEC3]
+            Target coordinates :math:`B=(x_1,y_1,z_1)` and next target coordinates
+            :math:`C=(x_2,y_2,z_2)`
+
+        Returns
+        -------
+            states: jdrones.data_models.States
+            reward: float
+            term: bool
+            trunc: bool
+            info: dict
+        """
+        A, B, C = np.array([self.env.env.state.pos, *action])
+        if np.allclose(B, C):
+            v_at_B = (0, 0, 0)
+        else:
+            v_at_B = self.calc_v_at_B(A, B, C, V=self.model.max_vel_ms * 0.3)
+        tgt_pos = B
+        tgt_vel = v_at_B
+        return self.env.step((tgt_pos, tgt_vel))

src/jdrones/types.py

@@ -19,3 +19,5 @@
 """:math:`(P_1,P_2,P_3,P_4)` propeller inputs"""
 PositionAction = VEC3
 """:math:`(x,y,z)` position inputs"""
+PositionVelocityAction = tuple[VEC3, VEC3]
+""":math:`(x,y,z)` position inputs, and :math:`(v_x,v_y,v_z)` velocity inputs"""

tests/conftest.py

@@ -9,6 +9,7 @@
 from jdrones.data_models import State
 from jdrones.data_models import URDFModel
 from jdrones.envs import FifthOrderPolyPositionDroneEnv
+from jdrones.envs import FifthOrderPolyPositionWithLookAheadDroneEnv
 from jdrones.envs import FirstOrderPolyPositionDroneEnv
 from jdrones.envs import LinearDynamicModelDroneEnv
 from jdrones.envs import LQRDroneEnv
@@ -212,10 +213,11 @@ def lqrdroneenv(env_default_kwargs):
     d.close()
 
 
-@pytest.fixture(params=[[[-0.1, 0.1], [-0.1, 0.1], [0, 0.2]]])
+@pytest.fixture(params=[[[-0.1, -0.1, 0], [0.1, 0.1, 0.2]]])
 def position_drone_action_space(request):
     a = np.array(request.param)
-    return spaces.Box(low=a[:, 0], high=a[:, 1])
+    low, high = a
+    return spaces.Box(low=low, high=high, shape=np.asarray(low).shape)
 
 
 def custom_position_action_space_wrapper(action_space, obj: type[BasePositionDroneEnv]):
@@ -243,3 +245,14 @@ def fifthorderpolypositiondroneenv(position_drone_action_space, env_default_kwar
     )(**env_default_kwargs)
     yield d
     d.close()
+
+
+@pytest.fixture
+def fifthorderpolypositionlookaheaddroneenv(
+    position_drone_action_space, env_default_kwargs
+):
+    d = custom_position_action_space_wrapper(
+        position_drone_action_space, FifthOrderPolyPositionWithLookAheadDroneEnv
+    )(**env_default_kwargs)
+    yield d
+    d.close()