[Refactor] DIff drive dynamics (#81)

* amend

* amend

* amend

* amend

* amend

* amend

* amend

* amend

* amend

Author: matteobettini

README.md

@@ -379,7 +379,7 @@ To create a fake screen you need to have `Xvfb` installed.
 | `line_trajectory.py`   | One agent is rewarded to move in a line trajectory.                                                                                                                                                                                                                                                                                                                                     | <img src="https://github.com/matteobettini/vmas-media/blob/main/media/scenarios/line_trajectory.gif?raw=true" alt="drawing" width="300"/>   |
 | `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"/>         |
+| `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"/> |
 
 ### [MPE](https://github.com/openai/multiagent-particle-envs)
 

vmas/scenarios/debug/diff_drive.py

@@ -5,7 +5,6 @@
 from typing import List
 
 import torch
-
 from vmas import render_interactively
 from vmas.simulator.core import Agent, World
 from vmas.simulator.dynamics.diff_drive import DiffDrive
@@ -25,11 +24,11 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs):
 
         The first agent has differential drive dynamics.
         You can control its forward input with the LEFT and RIGHT arrows.
-        You can control its rotation with N and M.
+        You can control its rotation with UP and DOWN.
 
         The second agent has standard vmas holonomic dynamics.
         You can control it with WASD
-        You can control its rotation withQ and E.
+        You can control its rotation with Q and E.
 
         """
         # T
@@ -46,7 +45,7 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs):
                     collide=True,
                     render_action=True,
                     u_range=[1, 1],
-                    u_multiplier=[1, 0.001],
+                    u_multiplier=[1, 1],
                     dynamics=DiffDrive(world, integration="rk4"),
                 )
             else:

vmas/simulator/dynamics/diff_drive.py

@@ -4,7 +4,6 @@
 import math
 
 import torch
-
 import vmas.simulator.core
 import vmas.simulator.utils
 from vmas.simulator.dynamics.common import Dynamics
@@ -23,35 +22,60 @@ def __init__(
         self.integration = integration
         self.world = world
 
-    def euler(self, f, rot):
-        return f(rot)
+    def euler(self, f, state):
+        return state + self.dt * f(state)
 
-    def runge_kutta(self, f, rot):
-        k1 = f(rot)
-        k2 = f(rot + self.dt * k1[2] / 2)
-        k3 = f(rot + self.dt * k2[2] / 2)
-        k4 = f(rot + self.dt * k3[2])
+    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 (1 / 6) * (k1 + 2 * k2 + 2 * k3 + k4)
+        return state + (self.dt / 6) * (k1 + 2 * k2 + 2 * k3 + k4)
 
     @property
     def needed_action_size(self) -> int:
         return 2
 
     def process_action(self):
-        u_forward = self.agent.action.u[:, 0]
-        u_rot = self.agent.action.u[:, 1]
+        velocity = self.agent.action.u[:, 0]  # Forward velocity
+        ang_velocity = self.agent.action.u[:, 1]  # Angular velocity
+
+        # Current state of the agent
+        state = torch.cat((self.agent.state.pos, self.agent.state.rot), dim=1)
 
-        def f(rot):
+        def f(state):
+            theta = state[:, 2]
             return torch.stack(
-                [u_forward * torch.cos(rot), u_forward * torch.sin(rot), u_rot], dim=0
+                [
+                    velocity * torch.cos(theta),
+                    velocity * torch.sin(theta),
+                    ang_velocity,
+                ],
+                dim=-1,
             )
 
         if self.integration == "euler":
-            u = self.euler(f, self.agent.state.rot.squeeze(-1))
+            new_state = self.euler(f, state)
         else:
-            u = self.runge_kutta(f, self.agent.state.rot.squeeze(-1))
+            new_state = self.runge_kutta(f, state)
+
+        # Calculate the change in state
+        delta_state = new_state - state
+
+        # Calculate the accelerations required to achieve the change in state
+        acceleration_x = delta_state[:, 0] / self.dt
+        acceleration_y = delta_state[:, 1] / self.dt
+        angular_acceleration = delta_state[:, 2] / 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 = self.agent.moment_of_inertia * angular_acceleration
 
-        self.agent.state.force[:, vmas.simulator.utils.X] = u[vmas.simulator.utils.X]
-        self.agent.state.force[:, vmas.simulator.utils.Y] = u[vmas.simulator.utils.Y]
-        self.agent.state.torque = u_rot.unsqueeze(-1)
+        # 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.unsqueeze(-1)

vmas/simulator/dynamics/kinematic_bicycle.py

@@ -5,7 +5,6 @@
 from typing import Union
 
 import torch
-
 import vmas.simulator.core
 import vmas.simulator.utils
 from vmas.simulator.dynamics.common import Dynamics
@@ -75,7 +74,7 @@ def f(state):
             dy = velocity * torch.sin(theta + beta)
             dtheta = velocity / (self.l_f + self.l_r) * torch.sin(beta)
             return torch.stack(
-                (dx, dy, dtheta), dim=1
+                (dx, dy, dtheta), dim=-1
             )  # Should return torch.Size([batch_size,3])
 
         # Select the integration method