refactor sim

furuta/logger.py

@@ -0,0 +1,57 @@
+import numpy as np
+import typing as tp
+import matplotlib.pyplot as plt
+
+from pathlib import Path
+
+STATE = ["phi", "theta", "phi_dot", "theta_dot"]
+
+class SimpleLogger:
+    def __init__(self):
+        self.times: tp.List[float] = []
+        self.states: tp.List[np.ndarray] = []
+
+    def update(self, time: float, state: np.ndarray):
+        self.times.append(time)
+        self.states.append(state)
+
+
+class Logger(SimpleLogger):
+    def __init__(self):
+        super().__init__()
+        self.controls: tp.List[float] = []
+        self.elapsed_times: tp.List[float] = []
+
+    def update(self, time: float, state: np.ndarray, control: float = 0.0, elapsed_time: float = 0.0):
+        super().update(time, state)
+        self.controls.append(control)
+        self.elapsed_times.append(elapsed_time)
+
+    def save(self, directory: Path):
+        np.save(directory / "times.npy", self.times)
+        np.save(directory / "states.npy", np.array(self.states))
+        np.save(directory / "controls.npy", self.controls)
+        np.save(directory / "elapsed_times.npy", self.elapsed_times)
+
+    def load(self, directory: Path):
+        self.states = np.load(directory / "states.npy")
+        self.times = np.load(directory / "times.npy")
+        self.controls = np.load(directory / "controls.npy")
+        self.elapsed_times = np.load(directory / "elapsed_times.npy")
+
+    def plot(self):
+        plt.figure(1)
+        for i in range(len(STATE)):
+            plt.subplot(2, 2, i + 1)
+            plt.plot(self.times, np.array(self.states)[:, i])
+            plt.title(STATE[i])
+
+        plt.figure(2)
+        plt.plot(self.times, self.controls)
+        plt.title("Control")
+
+        plt.figure(3)
+        plt.plot(self.times, self.elapsed_times)
+        plt.title("Elapsed Time")
+
+        plt.show()

furuta/sim.py

@@ -1,147 +1,28 @@
-import time
-from dataclasses import dataclass
-
-import matplotlib.pyplot as plt
 import numpy as np
 import pinocchio as pin
-from panda3d_viewer import Viewer
-from pinocchio.visualize.panda3d_visualizer import Panda3dVisualizer
 
 from furuta.robot import PendulumDynamics
 
-ROOT_DIR = "/home/pierfabre/pendulum_workspace/src/furuta/"
-STATE = ["phi", "theta", "phi_dot", "theta_dot"]
-
-
-@dataclass
-class RobotData:
-    time = 0.0
-    state = [0.0, 0.0, 0.0, 0.0]
-    control = 0.0
-    elapsed_time = 0.0
-
-    def update(self, time=None, state=None, control=None, elapsed_time=None):
-        if time is not None:
-            self.time = time
-        if state is not None:
-            self.state = state
-        if control is not None:
-            self.control = control
-        if elapsed_time is not None:
-            self.elapsed_time = elapsed_time
-
-
-class Logger:
-    def __init__(self):
-        self.times = []
-        self.states = []
-        self.controls = []
-        self.elapsed_times = []
-
-    def save(self, directory: str):
-        np.save(directory + "times.npy", self.times)
-        np.save(directory + "states.npy", self.states)
-        np.save(directory + "controls.npy", self.controls)
-        np.save(directory + "elapsed_times.npy", self.elapsed_times)
-
-    def update(self, data: RobotData):
-        self.times.append(data.time)
-        self.states.append(data.state)
-        self.controls.append(data.control)
-        self.elapsed_times.append(data.elapsed_time)
-
-    def load(self, directory: str):
-        self.states = np.load(directory + "states.npy")
-        self.times = np.load(directory + "times.npy")
-        self.control = np.load(directory + "controls.npy")
-        self.elapsed_times = np.load(directory + "elapsed_times.npy")
-
-    def plot(self):
-        plt.figure(1)
-        for i in range(len(STATE)):
-            plt.subplot(2, 2, i + 1)
-            plt.plot(self.times, np.array(self.states)[:, i])
-            plt.title(STATE[i])
-
-        plt.figure(2)
-        plt.plot(self.times, self.controls)
-        plt.title("Control")
-
-        plt.figure(3)
-        plt.plot(self.times, self.elapsed_times)
-        plt.title("Elapsed Time")
-
-        plt.show()
-
-
-class Integrator:
-    def __init__(self):
-        self.integrator = "euler"
-
-    # TODO : add rungekutta dopri integrator if needed
-
-    def euler(self, phi_ddot, th_ddot, state, dt):
-        phi, th, phi_dot, th_dot = state
-        phi_dot += phi_ddot * dt
-        th_dot += th_ddot * dt
-        phi += phi_dot * dt
-        th += th_dot * dt
-        return [phi, th, phi_dot, th_dot]
-
-    def __call__(self, phi_ddot, th_ddot, state, dt):
-        if self.integrator == "euler":
-            return self.euler(phi_ddot, th_ddot, state, dt)
-
 
 class SimulatedRobot:
-    def __init__(self, robot):
-        self.state = np.array([0.0, np.pi, 0.0, 0.0])
-        self.dt = 1e-6
-        self.dyn = PendulumDynamics()
-        self.integrator = Integrator()
+    def __init__(self, robot: pin.RobotWrapper, init_state: np.ndarray, dt: float):
         self.robot = robot
-        self.model = self.robot.model
-        self.data = self.model.createData()
+        self.state = init_state
+        self.dt = dt
+        self.dyn = PendulumDynamics(robot)
 
-    def step(self, u: float, dt: float):
+    def step(self, u: float, dt: float) -> np.ndarray:
         """Simulate the robot for a given control input for a given duration :param u: control
         input :param dt: duration of the control input."""
         t = 0.0
         q = np.array(self.state[:2])
         v = np.array(self.state[2:])
         while t < dt:
-            tau = np.array([u, 0.0])
-            tau[1] = -self.dyn.f_p * v[1]
-            tau[0] -= self.dyn.motor.f_v * v[0] + self.dyn.motor.f_c * np.sign(v[0])
-            a = pin.aba(self.model, self.data, q, v, tau)
+            # Compute dynamics
+            a = self.dyn(q, v, u)
+            # Integrate
             v += a * self.dt
-            q = pin.integrate(self.model, q, v * self.dt)
+            q = pin.integrate(self.robot.model, q, v * self.dt)
             t += self.dt
         self.state = np.concatenate((q, v))
         return self.state
-
-
-class RobotViewer:
-    def __init__(self, robot: pin.RobotWrapper = None):
-        self.robot = robot
-        viewer = Viewer()
-        # Attach the robot to the viewer scene
-        self.robot.setVisualizer(Panda3dVisualizer())
-        self.robot.initViewer(viewer=viewer)
-        self.robot.loadViewerModel(group_name=self.robot.model.name)
-
-    def display(self, q):
-        self.robot.display(q)
-
-    def animate(self, times, states):
-        # Initial state
-        q = np.array(states)[:, :2]
-        self.display(q[0])
-        time.sleep(1.0)
-        tic = time.time()
-        for i in range(1, len(times)):
-            toc = time.time()
-            time.sleep(max(0, times[i] - times[i - 1] - (toc - tic)))
-            tic = time.time()
-            if i % 10 == 0:
-                self.display(q[i])

furuta/utils.py

@@ -1,5 +1,8 @@
 # https://git.ias.informatik.tu-darmstadt.de/quanser/clients/-/blob/v0.1.1/quanser_robots/common.py
+from pathlib import Path
+
 import numpy as np
+import pinocchio as pin
 from gymnasium import spaces
 from scipy import signal
 
@@ -9,6 +12,14 @@
 ALPHA_DOT = 3
 
 
+ROOT_DIR = Path.home() / "Documents/perso/pendulum_workspace/"
+
+robot = pin.RobotWrapper.BuildFromURDF(
+    str(ROOT_DIR / "src/robot/hardware/furuta.urdf"),
+    package_dirs=[str(ROOT_DIR / "src/robot/hardware/CAD/stl/")],
+)
+
+
 class VelocityFilter:
     """Discrete velocity filter derived from a continuous one."""
 

furuta/viewer.py

@@ -0,0 +1,37 @@
+import time
+
+import numpy as np
+import pinocchio as pin
+
+from panda3d_viewer import Viewer
+from pinocchio.visualize.panda3d_visualizer import Panda3dVisualizer
+from furuta.logger import SimpleLogger
+
+
+class RobotViewer:
+    def __init__(self, robot: pin.RobotWrapper = None):
+        self.robot = robot
+        viewer = Viewer()
+        # Attach the robot to the viewer scene
+        self.robot.setVisualizer(Panda3dVisualizer())
+        self.robot.initViewer(viewer=viewer)
+        self.robot.loadViewerModel(group_name=self.robot.model.name)
+
+    def display(self, q: np.ndarray):
+        self.robot.display(q)
+
+    def animate(self, times: np.ndarray, states: np.ndarray):
+        # Initial state
+        q = np.array(states)[:, :2]
+        self.display(q[0])
+        time.sleep(1.0)
+        tic = time.time()
+        for i in range(1, len(times)):
+            toc = time.time()
+            time.sleep(max(0, times[i] - times[i - 1] - (toc - tic)))
+            tic = time.time()
+            if i % 10 == 0:
+                self.display(q[i])
+
+    def animate_log(self, log: SimpleLogger):
+        self.animate(log.times, log.states)