From 30f9934154aa2700fc6b7cab9175598329c80548 Mon Sep 17 00:00:00 2001
From: Michel Breyer <michel.p.breyer@gmail.com>
Date: Tue, 16 Mar 2021 17:07:33 +0100
Subject: [PATCH] Add control loop

---
 controllers.py | 12 ++++++++----
 robot_sim.py   | 19 ++++++++++++-------
 run_demo.py    | 17 ++++-------------
 3 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/controllers.py b/controllers.py
index 7202433..632bb19 100644
--- a/controllers.py
+++ b/controllers.py
@@ -2,15 +2,19 @@ import numpy as np
 
 
 class CartesianPoseController:
-    def __init__(self, model, x0):
+    def __init__(self, robot, model, rate):
+        self.robot = robot
         self.model = model
-        self.x_d = x0
+        self.rate = rate
+        self.x_d = None
         self.kp = np.ones(6) * 5.0
 
     def set_target(self, pose):
         self.x_d = pose
 
-    def update(self, q, dq):
+    def update(self):
+        q, _ = self.robot.get_state()
+
         x = self.model.pose(q)
         x_d = self.x_d
 
@@ -22,4 +26,4 @@ class CartesianPoseController:
         J_pinv = np.linalg.pinv(J)
         cmd = np.dot(J_pinv, self.kp * err)
 
-        return cmd
+        self.robot.set_desired_joint_velocities(cmd)
diff --git a/robot_sim.py b/robot_sim.py
index 923fe8d..b522f81 100644
--- a/robot_sim.py
+++ b/robot_sim.py
@@ -6,6 +6,7 @@ import pybullet_data
 
 from sensor_msgs.msg import JointState
 
+from model import *
 from utils import *
 
 
@@ -107,13 +108,14 @@ class Camera:
 
 
 class SimPandaEnv(object):
-    def __init__(self, gui, dt=1.0 / 240.0, publish_state=True):
+    def __init__(self, gui, publish_state=True):
         self.gui = gui
-        self.dt = dt
+        self.rate = 240
         p.connect(p.GUI if gui else p.DIRECT)
         p.setAdditionalSearchPath(pybullet_data.getDataPath())
-        p.setTimeStep(dt)
+        p.setTimeStep(1.0 / self.rate)
         p.setGravity(0.0, 0.0, -9.8)
+        self.step_cnt = 0
 
         p.loadURDF("plane.urdf")
         p.loadURDF(
@@ -127,6 +129,7 @@ class SimPandaEnv(object):
         self.arm.load(Transform(Rotation.identity(), np.r_[-0.8, 0.0, 0.4]))
         self.gripper = PandaGripper()
         self.gripper.uid = self.arm.uid
+        self.model = Model("panda_link0", "panda_ee")
         self.camera = Camera(self.arm.uid, 12, 320, 240, 1.047, 0.2, 2.0)
 
         if publish_state:
@@ -134,13 +137,15 @@ class SimPandaEnv(object):
             rospy.Timer(rospy.Duration(1.0 / 30), self._publish_state)
 
     def forward(self, dt):
-        steps = int(dt / self.dt)
+        steps = int(dt * self.rate)
         for _ in range(steps):
-            self.step()
+            self._step()
 
-    def step(self):
+    def _step(self):
+        if self.step_cnt % int(self.rate / self.controller.rate) == 0:
+            self.controller.update()
         p.stepSimulation()
-        time.sleep(self.dt)
+        time.sleep(1.0 / self.rate)
 
     def _publish_state(self, event):
         positions, velocities = self.arm.get_state()
diff --git a/run_demo.py b/run_demo.py
index 10510e1..d021bf1 100644
--- a/run_demo.py
+++ b/run_demo.py
@@ -1,30 +1,21 @@
 import rospy
 
 from controllers import *
-from model import *
 from robot_sim import *
 from utils import *
 
 
 gui = True
-dt = 1.0 / 60.0
 
 rospy.init_node("demo")
 
 env = SimPandaEnv(gui)
-model = Model("panda_link0", "panda_ee")
+env.controller = CartesianPoseController(env.arm, env.model, 60)
 
 q, dq = env.arm.get_state()
-x0 = model.pose(q)
-
-controller = CartesianPoseController(model, x0)
-
+x0 = env.model.pose(q)
 marker = InteractiveMarkerWrapper("target", "panda_link0", x0)
 
 while True:
-    controller.set_target(marker.pose)
-    q, dq = env.arm.get_state()
-    cmd = controller.update(q, dq)
-    env.arm.set_desired_joint_velocities(cmd)
-
-    env.forward(dt)
+    env.controller.set_target(marker.pose)
+    env.forward(0.1)