Merge pull request #399 from IMRCLab/issue398

sim - backend - add neuralswarm

crazyflie_examples/crazyflie_examples/swap.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+
+from crazyflie_py import Crazyswarm
+import numpy as np
+
+
+def main():
+    Id2 = 231
+    Id1 = 5
+    Pos1 = np.array([0.0, -0.2, 0.0])
+    Pos2 = np.array([0.0, 0.2, 0.0])
+    Height1 = 0.4
+    Height2 = 0.5
+    swapTime = 3
+
+    swarm = Crazyswarm()
+    timeHelper = swarm.timeHelper
+    allcfs = swarm.allcfs
+
+    allcfs.takeoff(targetHeight=Height1, duration=3.0)
+    timeHelper.sleep(3.5)
+
+    # go to initial positions
+    allcfs.crazyfliesById[Id1].goTo(Pos1 + np.array([0, 0, Height1]), 0, 3.0)
+    allcfs.crazyfliesById[Id2].goTo(Pos2 + np.array([0, 0, Height2]), 0, 3.0)
+    timeHelper.sleep(3.5)
+
+    # swap 1
+    allcfs.crazyfliesById[Id1].goTo(Pos2 + np.array([0, 0, Height1]), 0, swapTime)
+    allcfs.crazyfliesById[Id2].goTo(Pos1 + np.array([0, 0, Height2]), 0, swapTime)
+    timeHelper.sleep(swapTime + 1.5)
+
+    # swap 2
+    allcfs.crazyfliesById[Id1].goTo(Pos1 + np.array([0, 0, Height1]), 0, swapTime)
+    allcfs.crazyfliesById[Id2].goTo(Pos2 + np.array([0, 0, Height2]), 0, swapTime)
+    timeHelper.sleep(swapTime + 1.5)
+
+    allcfs.land(targetHeight=0.02, duration=3.0)
+    timeHelper.sleep(3.5)
+
+
+if __name__ == '__main__':
+    main()

crazyflie_examples/setup.cfg

@@ -7,3 +7,4 @@ console_scripts =
     multi_trajectory = crazyflie_examples.multi_trajectory:main
     cmd_full_state = crazyflie_examples.cmd_full_state:main
     set_param = crazyflie_examples.set_param:main
+    swap = crazyflie_examples.swap:main

crazyflie_sim/crazyflie_sim/backend/neuralswarm.py

@@ -0,0 +1,170 @@
+"""
+This implementes interaction force prediction using NeuralSwarm(2).
+
+See https://github.com/aerorobotics/neural-swarm
+
+Logic copied from https://github.com/aerorobotics/neural-swarm/blob/master/planning/robots.py
+"""
+
+from pathlib import Path
+
+import numpy as np
+from rclpy.node import Node
+from rclpy.time import Time
+from rosgraph_msgs.msg import Clock
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .np import Quadrotor
+from ..sim_data_types import Action, State
+
+
+# H is the dimension of the hidden state
+class phi_Net(nn.Module):
+
+    def __init__(self, inputdim=6, hiddendim=40):
+        super(phi_Net, self).__init__()
+        self.fc1 = nn.Linear(inputdim, 25)
+        self.fc2 = nn.Linear(25, 40)
+        self.fc3 = nn.Linear(40, 40)
+        self.fc4 = nn.Linear(40, hiddendim)
+
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = F.relu(self.fc3(x))
+        x = self.fc4(x)
+        return x
+
+
+class rho_Net(nn.Module):
+
+    def __init__(self, hiddendim=40):
+        super(rho_Net, self).__init__()
+        self.fc1 = nn.Linear(hiddendim, 40)
+        self.fc2 = nn.Linear(40, 40)
+        self.fc3 = nn.Linear(40, 40)
+        self.fc4 = nn.Linear(40, 1)
+
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = F.relu(self.fc3(x))
+        x = self.fc4(x)
+        return x
+
+
+class NeuralSwarm:
+
+    def __init__(self, model_folder):
+        self.H = 20
+        self.rho_L_net = rho_Net(hiddendim=self.H)
+        self.phi_L_net = phi_Net(inputdim=6, hiddendim=self.H)  # x,y,z,vx,vy,vz
+        self.rho_L_net.load_state_dict(torch.load('{}/rho_L.pth'.format(model_folder)))
+        self.phi_L_net.load_state_dict(torch.load('{}/phi_L.pth'.format(model_folder)))
+        self.rho_S_net = rho_Net(hiddendim=self.H)
+        self.phi_S_net = phi_Net(inputdim=6, hiddendim=self.H)  # x,y,z,vx,vy,vz
+        self.rho_S_net.load_state_dict(torch.load('{}/rho_S.pth'.format(model_folder)))
+        self.phi_S_net.load_state_dict(torch.load('{}/phi_S.pth'.format(model_folder)))
+        self.phi_G_net = phi_Net(inputdim=4, hiddendim=self.H)  # z,vx,vy,vz
+        self.phi_G_net.load_state_dict(torch.load('{}/phi_G.pth'.format(model_folder)))
+
+    def compute_Fa(self, data_self, data_neighbors):
+        rho_input = torch.zeros(self.H)
+        cftype, x = data_self
+        for cftype_neighbor, x_neighbor in data_neighbors:
+            x_12 = torch.zeros(6)
+            x_12 = (x_neighbor - x).float()
+            if abs(x_12[0]) < 0.2 and abs(x_12[1]) < 0.2 and abs(x_12[3]) < 1.5:
+                if cftype_neighbor == 'small' or cftype_neighbor == 'small_powerful_motors':
+                    rho_input += self.phi_S_net(x_12)
+                elif cftype_neighbor == 'large':
+                    rho_input += self.phi_L_net(x_12)
+                else:
+                    raise Exception('Unknown cftype!')
+
+        # interaction with the ground
+        x_12 = torch.zeros(4)
+        x_12[0] = 0 - x[2]
+        x_12[1:4] = -x[3:6]
+        rho_input += self.phi_G_net(x_12)
+
+        if cftype == 'small' or cftype == 'small_powerful_motors':
+            faz = self.rho_S_net(rho_input)
+        elif cftype == 'large':
+            faz = self.rho_L_net(rho_input)
+        else:
+            raise Exception('Unknown cftype!')
+
+        return np.array([0, 0, faz[0].item()])
+
+
+class Backend:
+    """Backend that is based on the one defined in np.py."""
+
+    def __init__(self, node: Node, names: list[str], states: list[State]):
+        self.node = node
+        self.names = names
+        self.clock_publisher = node.create_publisher(Clock, 'clock', 10)
+        self.t = 0
+        self.dt = 0.0005
+
+        self.uavs = []
+        for state in states:
+            uav = Quadrotor(state)
+            self.uavs.append(uav)
+        self.neuralswarm = NeuralSwarm(Path(__file__).parent / 'data/neuralswarm2')
+
+    def time(self) -> float:
+        return self.t
+
+    def step(self, states_desired: list[State], actions: list[Action]) -> list[State]:
+        # advance the time
+        self.t += self.dt
+
+        fa_data = []
+        for uav in self.uavs:
+            fa_data.append(('small', torch.hstack(
+                (torch.tensor(uav.state.pos), torch.tensor(uav.state.vel)))))
+
+        next_states = []
+        for k, (uav, action) in enumerate(zip(self.uavs, actions)):
+            # estimate F_a
+            # print(k, fa_data[k], fa_data[0:k] + fa_data[k+1:])
+            f_a = self.neuralswarm.compute_Fa(fa_data[k], fa_data[0:k] + fa_data[k+1:])
+            # convert grams to Newtons
+            f_a = f_a / 1000 * 9.81
+            # print(f_a)
+            uav.step(action, self.dt, f_a)
+            next_states.append(uav.state)
+
+        # print(states_desired, actions, next_states)
+        # publish the current clock
+        clock_message = Clock()
+        clock_message.clock = Time(seconds=self.time()).to_msg()
+        self.clock_publisher.publish(clock_message)
+
+        return next_states
+
+    def shutdown(self):
+        pass
+
+
+if __name__ == '__main__':
+
+    # test case, see Fig 5g in NeuralSwarm2 paper
+    ns = NeuralSwarm(Path(__file__).parent / 'data/neuralswarm2')
+
+    # x, y, z, vx, vy, vz
+    states = torch.tensor([
+        [0, 0, 0.6, 0, 0, 0],
+        [0, -0.1, 0.5, 0, 0, 0],
+        [0, 0.1, 0.5, 0, 0, 0],
+        [0, 0, 0.3, 0, 0, 0],
+    ])
+
+    print(ns.compute_Fa(('small', states[3]),
+                        [('small', states[0]),
+                        ('small', states[1]),
+                        ('small', states[2])]))

crazyflie_sim/crazyflie_sim/backend/np.py

@@ -81,7 +81,7 @@ def __init__(self, state):
 
         self.state = state
 
-    def step(self, action, dt):
+    def step(self, action, dt, f_a=np.zeros(3)):
 
         # convert RPM -> Force
         def rpm_to_force(rpm):
@@ -101,10 +101,10 @@ def rpm_to_force(rpm):
         # dynamics
         # dot{p} = v
         pos_next = self.state.pos + self.state.vel * dt
-        # mv = mg + R f_u
+        # mv = mg + R f_u + f_a
         vel_next = self.state.vel + (
             np.array([0, 0, -self.g]) +
-            rowan.rotate(self.state.quat, f_u) / self.mass) * dt
+            (rowan.rotate(self.state.quat, f_u) + f_a) / self.mass) * dt
 
         # dot{R} = R S(w)
         # to integrate the dynamics, see