[Dynamics] Drone dynamics (#83)

gy2256 · matteobettini · web-flow · commit d58854851406 · 2024-02-16T14:41:18.000Z
* Add drone dynamics

* VMAS-1.4.0

* Update vmas/simulator/dynamics/drone.py

Co-authored-by: Matteo Bettini &lt;55539777+matteobettini@users.noreply.github.com&gt;

* Update drone.py

1. the mass is obtained from self.agent.mass now
2. Updated how the drone_state is initialized

* Update torque (action input) with correct indexing

* implemented reset function

* Update drone_state dimension

* removed vmas_state

* Bug fix

The agent now behaves correctly using the interactive rendering.

* Adding need_reset function

* state is unpacked correctly now in def f

* Amend

* empty

* Amend

* Amend

---------

Co-authored-by: Matteo Bettini &lt;mb2389@cl.cam.ac.uk&gt;
Co-authored-by: Matteo Bettini &lt;55539777+matteobettini@users.noreply.github.com&gt;
diff --git a/README.md b/README.md
@@ -246,7 +246,7 @@ customizable. Examples are: drag, friction, gravity, simulation timestep, non-di
 - **Agent actions**: Agents' physical actions are 2D forces for holonomic motion. Agent rotation can also be controlled through a torque action (activated by setting `agent.action.u_rot_range` at agent creation time). Agents can also be equipped with continuous or discrete communication actions.
 - **Action preprocessing**: By implementing the `process_action` function of a scenario, you can modify the agents' actions before they are passed to the simulator. This is used in `controllers` (where we provide different types of controllers to use) and `dynamics` (where we provide custom robot dynamic models).
 - **Controllers**: Controllers are components that can be appended to the neural network policy or replace it completely.  We provide a `VelocityController` which can be used to treat input actions as velocities (instead of default vmas input forces). This PID controller takes velocities and outputs the forces which are fed to the simulator. See the `vel_control` debug scenario for an example.
-- **Dynamic models**: VMAS simulates holonomic dynamics models by default. Custom dynamics can be chosen at agent creation time. Implementations now include `DiffDriveDynamics` for differential drive robots and `KinematicBicycleDynamics` for kinematic bicycle model. See `diff_drive` and `kinematic_bicycle` debug scenarios for examples.
+- **Dynamic models**: VMAS simulates holonomic dynamics models by default. Custom dynamics can be chosen at agent creation time. Implementations now include `DiffDriveDynamics` for differential drive robots, `KinematicBicycleDynamics` for kinematic bicycle model, and `Drone` for quadcopter dynamics. See `diff_drive`, `kinematic_bicycle` and `drone` debug scenarios for examples.
 - **Differentiable**: By setting `grad_enabled=True` when creating an environment, the simulator will be differentiable, allowing gradients flowing through any of its function.
 
 ## Creating a new scenario
@@ -380,6 +380,7 @@ To create a fake screen you need to have `Xvfb` installed.
 | `circle_trajectory.py` | One agent is rewarded to move in a circle trajectory at the `desired_radius`.                                                                                                                                                                                                                                                                                                           | <img src="https://github.com/matteobettini/vmas-media/blob/main/media/scenarios/circle_trajectory.gif?raw=true" alt="drawing" width="300"/> |
 | `diff_drive.py`        | An example of the `diff_drive` dynamic model constraint. Both agents have rotational actions which can be controlled interactively.  The first agent has differential drive dynamics. The second agent has standard vmas holonomic dynamics.                                                                                                                                            | <img src="https://github.com/matteobettini/vmas-media/blob/main/media/scenarios/diff_drive.gif?raw=true" alt="drawing" width="300"/>        |
 | `kinematic_bicycle.py` | An example of `kinematic_bicycle` dynamic model constraint. Both agents have rotational actions which can be controlled interactively.  The first agent has kinematic bicycle model dynamics. The second agent has standard vmas holonomic dynamics.                                                                                                                                    | <img src="https://github.com/matteobettini/vmas-media/blob/main/media/scenarios/kinematic_bicycle.gif?raw=true" alt="drawing" width="300"/> |
+| `drone.py`             | An example of the `drone` dynamic model.                                                                                                                                                                                                                                                                                                                                                | <img src="https://github.com/matteobettini/vmas-media/blob/main/media/scenarios/drone.gif?raw=true" alt="drawing" width="300"/>             |
 
 ### [MPE](https://github.com/openai/multiagent-particle-envs)
 
@@ -408,7 +409,7 @@ To create a fake screen you need to have `Xvfb` installed.
 - [ ] Improve test efficiency and add new tests
 - [ ] Implement 1D camera sensor
 - [ ] Implement 2D birds eye view camera sensor
-- [ ] Implement 2D drone dynamics
+- [X] Implement 2D drone dynamics
 - [X] Allow any number of actions
 - [X] Improve VMAS performance
 - [X] Dict obs support in torchrl
diff --git a/vmas/interactive_rendering.py b/vmas/interactive_rendering.py
@@ -31,8 +31,8 @@ class InteractiveEnv:
 
     You can change agent by pressing TAB
     You can reset the environment by pressing R
-    You can move agents with the arrow keys and if the agent has a rotational action you can control it with M, N
-    If you have more than 1 agent, you can control another one with W,A,S,D and Q,E for eventual rotational actions
+    You can control agent actions with the arrow keys and M/N (left/right control the first action, up/down control the second, M/N controls the third)
+    If you have more than 1 agent, you can control another one with W,A,S,D and Q,E in the same way.
     and switch the agent with these controls using LSHIFT
     """
 
@@ -305,8 +305,8 @@ def render_interactively(
 
     You can change agent by pressing TAB
     You can reset the environment by pressing R
-    You can move agents with the arrow keys and if the agent has a rotational action you can control it with M, N
-    If you have more than 1 agent, you can control another one with W,A,S,D and Q,E for eventual rotational actions
+    You can control agent actions with the arrow keys and M/N (left/right control the first action, up/down control the second, M/N controls the third)
+    If you have more than 1 agent, you can control another one with W,A,S,D and Q,E in the same way.
     and switch the agent with these controls using LSHIFT
     """
 
@@ -333,8 +333,8 @@ def render_interactively(
     #
     # You can change agent by pressing TAB
     # You can reset the environment by pressing R
-    # You can move agents with the arrow keys and if the agent has a rotational action you can control it with M, N
-    # If you have more than 1 agent, you can control another one with W,A,S,D and Q,E for eventual rotational actions
+    # You can control agent actions with the arrow keys and M/N (left/right control the first action, up/down control the second, M/N controls the third)
+    # If you have more than 1 agent, you can control another one with W,A,S,D and Q,E in the same way.
     # and switch the agent with these controls using LSHIFT
 
     scenario_name = "waterfall"
diff --git a/vmas/scenarios/debug/drone.py b/vmas/scenarios/debug/drone.py
@@ -0,0 +1,113 @@
+#  Copyright (c) 2024.
+#  ProrokLab (https://www.proroklab.org/)
+#  All rights reserved.
+
+import typing
+from typing import List
+
+import torch
+
+from vmas import render_interactively
+from vmas.simulator.core import Agent, World
+from vmas.simulator.dynamics.drone import Drone
+from vmas.simulator.scenario import BaseScenario
+from vmas.simulator.utils import Color, ScenarioUtils
+
+if typing.TYPE_CHECKING:
+    from vmas.simulator.rendering import Geom
+
+
+class Scenario(BaseScenario):
+    def make_world(self, batch_dim: int, device: torch.device, **kwargs):
+        """
+        Drone example scenario
+        Run this file to try it out.
+
+        You can control the three input torques using left/right arrows, up/down arrows, and m/n.
+        """
+        self.plot_grid = True
+        self.n_agents = kwargs.get("n_agents", 2)
+
+        # Make world
+        world = World(batch_dim, device, substeps=10)
+
+        for i in range(self.n_agents):
+            agent = Agent(
+                name=f"drone_{i}",
+                collide=True,
+                render_action=True,
+                u_range=[0.00001, 0.00001, 0.00001],  # torque_x, torque_y, torque_z
+                u_multiplier=[1, 1, 1],
+                action_size=3,  # We feed only the torque actions to interactively control the drone in the debug scenario
+                # In non-debug cases, remove this line and the `process_action` function in this file
+                dynamics=Drone(world, integration="rk4"),
+            )
+            world.add_agent(agent)
+
+        return world
+
+    def reset_world_at(self, env_index: int = None):
+        ScenarioUtils.spawn_entities_randomly(
+            self.world.agents,
+            self.world,
+            env_index,
+            min_dist_between_entities=0.1,
+            x_bounds=(-1, 1),
+            y_bounds=(-1, 1),
+        )
+
+    def reward(self, agent: Agent):
+        return torch.zeros(self.world.batch_dim, device=self.world.device)
+
+    def process_action(self, agent: Agent):
+        torque = agent.action.u
+        thrust = torch.full(
+            (self.world.batch_dim, 1),
+            agent.mass * agent.dynamics.g,
+            device=self.world.device,
+        )  # Add a fixed thrust to make sure the agent is not falling
+        agent.action.u = torch.cat([thrust, torque], dim=-1)
+
+    def observation(self, agent: Agent):
+        observations = [
+            agent.state.pos,
+            agent.state.vel,
+        ]
+        return torch.cat(
+            observations,
+            dim=-1,
+        )
+
+    def done(self):
+        return torch.any(
+            torch.stack(
+                [agent.dynamics.needs_reset() for agent in self.world.agents], dim=-1
+            ),
+            dim=-1,
+        )
+
+    def extra_render(self, env_index: int = 0) -> "List[Geom]":
+        from vmas.simulator import rendering
+
+        geoms: List[Geom] = []
+
+        # Agent rotation
+        for agent in self.world.agents:
+            color = Color.BLACK.value
+            line = rendering.Line(
+                (0, 0),
+                (0.1, 0),
+                width=1,
+            )
+            xform = rendering.Transform()
+            xform.set_rotation(agent.state.rot[env_index])
+            xform.set_translation(*agent.state.pos[env_index])
+            line.add_attr(xform)
+            line.set_color(*color)
+            geoms.append(line)
+
+        return geoms
+
+
+if __name__ == "__main__":
+    render_interactively(__file__, control_two_agents=True)
diff --git a/vmas/simulator/dynamics/common.py b/vmas/simulator/dynamics/common.py
@@ -29,11 +29,15 @@ def agent(self):
     def agent(self, value):
         if self._agent is not None:
             raise ValueError("Agent in dynamics has already been set")
-        if value.action_size < self.needed_action_size:
+        self._agent = value
+
+    def check_and_process_action(self):
+        action = self.agent.action.u
+        if action.shape[1] < self.needed_action_size:
             raise ValueError(
-                f"Agent action size {value.action_size} is less than the required dynamics action size {self.needed_action_size}"
+                f"Agent action size {action.shape[1]} is less than the required dynamics action size {self.needed_action_size}"
             )
-        self._agent = value
+        self.process_action()
 
     @property
     @abc.abstractmethod
diff --git a/vmas/simulator/dynamics/drone.py b/vmas/simulator/dynamics/drone.py
@@ -0,0 +1,156 @@
+#  Copyright (c) 2024.
+#  ProrokLab (https://www.proroklab.org/)
+#  All rights reserved.
+
+from typing import Union
+
+import torch
+from torch import Tensor
+
+import vmas.simulator.core
+import vmas.simulator.utils
+from vmas.simulator.dynamics.common import Dynamics
+
+
+class Drone(Dynamics):
+    def __init__(
+        self,
+        world: vmas.simulator.core.World,
+        I_xx: float = 8.1e-3,
+        I_yy: float = 8.1e-3,
+        I_zz: float = 14.2e-3,
+        integration: str = "rk4",
+    ):
+        super().__init__()
+
+        assert integration in (
+            "rk4",
+            "euler",
+        )
+
+        self.integration = integration
+        self.I_xx = I_xx
+        self.I_yy = I_yy
+        self.I_zz = I_zz
+        self.world = world
+        self.g = 9.81
+        self.dt = world.dt
+        self.reset()
+
+    def reset(self, index: Union[Tensor, int] = None):
+        if index is None:
+            # Drone state: phi(roll), theta (pitch), psi (yaw),
+            #              p (roll_rate), q (pitch_rate), r (yaw_rate),
+            #              x_dot (vel_x), y_dot (vel_y), z_dot (vel_z),
+            #              x (pos_x), y (pos_y), z (pos_z)
+            self.drone_state = torch.zeros(
+                self.world.batch_dim,
+                12,
+                device=self.world.device,
+            )
+        else:
+            self.drone_state[index] = 0.0
+
+    def needs_reset(self) -> Tensor:
+        # Constraint roll and pitch within +-30 degrees
+        return torch.any(self.drone_state[:, :2].abs() > 30 * (torch.pi / 180), dim=-1)
+
+    def euler(self, f, state):
+        return state + self.dt * f(state)
+
+    def runge_kutta(self, f, state):
+        k1 = f(state)
+        k2 = f(state + self.dt * k1 / 2)
+        k3 = f(state + self.dt * k2 / 2)
+        k4 = f(state + self.dt * k3)
+        return state + (self.dt / 6) * (k1 + 2 * k2 + 2 * k3 + k4)
+
+    @property
+    def needed_action_size(self) -> int:
+        return 4
+
+    def process_action(self):
+        u = self.agent.action.u
+        thrust = u[:, 0]  # Thrust, sum of all propeller thrusts
+        torque = u[:, 1:4]  # Torque in x, y, z direction
+
+        thrust += self.agent.mass * self.g  # Ensure the drone is not falling
+
+        self.drone_state[:, 9] = self.agent.state.pos[:, 0]  # x
+        self.drone_state[:, 10] = self.agent.state.pos[:, 1]  # y
+        self.drone_state[:, 2] = self.agent.state.rot[:, 0]  # psi (yaw)
+
+        def f(state):
+            phi = state[:, 0]
+            theta = state[:, 1]
+            psi = state[:, 2]
+            p = state[:, 3]
+            q = state[:, 4]
+            r = state[:, 5]
+            x_dot = state[:, 6]
+            y_dot = state[:, 7]
+            z_dot = state[:, 8]
+
+            c_phi = torch.cos(phi)
+            s_phi = torch.sin(phi)
+            c_theta = torch.cos(theta)
+            s_theta = torch.sin(theta)
+            c_psi = torch.cos(psi)
+            s_psi = torch.sin(psi)
+
+            # Postion Dynamics
+            x_ddot = (
+                (c_phi * s_theta * c_psi + s_phi * s_psi) * thrust / self.agent.mass
+            )
+            y_ddot = (
+                (c_phi * s_theta * s_psi - s_phi * c_psi) * thrust / self.agent.mass
+            )
+            z_ddot = (c_phi * c_theta) * thrust / self.agent.mass - self.g
+            # Angular velocity dynamics
+            p_dot = (torque[:, 0] - (self.I_yy - self.I_zz) * q * r) / self.I_xx
+            q_dot = (torque[:, 1] - (self.I_zz - self.I_xx) * p * r) / self.I_yy
+            r_dot = (torque[:, 2] - (self.I_xx - self.I_yy) * p * q) / self.I_zz
+
+            return torch.stack(
+                [
+                    p,
+                    q,
+                    r,
+                    p_dot,
+                    q_dot,
+                    r_dot,
+                    x_ddot,
+                    y_ddot,
+                    z_ddot,
+                    x_dot,
+                    y_dot,
+                    z_dot,
+                ],
+                dim=-1,
+            )
+
+        if self.integration == "euler":
+            new_drone_state = self.euler(f, self.drone_state)
+        else:
+            new_drone_state = self.runge_kutta(f, self.drone_state)
+
+        # Calculate the change in state
+        delta_state = new_drone_state - self.drone_state
+        self.drone_state = new_drone_state
+
+        # Calculate the accelerations required to achieve the change in state
+        acceleration_x = delta_state[:, 6] / self.dt
+        acceleration_y = delta_state[:, 7] / self.dt
+        angular_acceleration = delta_state[:, 5] / self.dt
+
+        # Calculate the forces required for the linear accelerations
+        force_x = self.agent.mass * acceleration_x
+        force_y = self.agent.mass * acceleration_y
+
+        # Calculate the torque required for the angular acceleration
+        torque_yaw = self.agent.moment_of_inertia * angular_acceleration
+
+        # Update the physical force and torque required for the user inputs
+        self.agent.state.force[:, vmas.simulator.utils.X] = force_x
+        self.agent.state.force[:, vmas.simulator.utils.Y] = force_y
+        self.agent.state.torque = torque_yaw.unsqueeze(-1)
diff --git a/vmas/simulator/scenario.py b/vmas/simulator/scenario.py
@@ -65,7 +65,7 @@ def env_process_action(self, agent: Agent):
             agent.action_callback(self.world)
         # Customizable action processor
         self.process_action(agent)
-        agent.dynamics.process_action()
+        agent.dynamics.check_and_process_action()
 
     @abstractmethod
     def make_world(self, batch_dim: int, device: torch.device, **kwargs) -> World:
