Tested TP-RMP with RMPflow

Signed-off-by: An Thai Le <an.thai.le97@gmail.com>
ASUDCT · Aug 26, 2021 · c30c598 · c30c598
1 parent 491facd
commit c30c598
Show file tree

Hide file tree

Showing 11 changed files with 237 additions and 38 deletions.
diff --git a/data/assets/sphere/sphere-template.urdf b/data/assets/sphere/sphere-template.urdf
@@ -0,0 +1,24 @@
+<?xml version="1.0" ?>
+<robot name="simple_sphere">
+  <link name="sphere">
+    <inertial>
+      <mass value="MASS0" />
+      <origin xyz="0 0 0" /> 
+      <inertia  ixx="1.0" ixy="0.0"  ixz="0.0"  iyy="1.0"  iyz="0.0"  izz="1.0" />
+    </inertial>
+    <visual>
+      <!-- visual origin is defined w.r.t. link local coordinate system -->
+      <origin xyz="0 0 0" rpy="0 0 0" />
+      <geometry>
+        <sphere radius="RADIUS0"/>
+      </geometry>
+    </visual>
+    <collision>
+      <!-- collision origin is defined w.r.t. link local coordinate system -->
+      <origin xyz="0 0 0" rpy="0 0 0" />
+      <geometry>
+        <sphere radius="RADIUS0"/>
+      </geometry>
+    </collision>
+  </link>
+</robot>
diff --git a/data/assets/sphere/sphere.urdf b/data/assets/sphere/sphere.urdf
@@ -0,0 +1,24 @@
+<?xml version="1.0" ?>
+<robot name="simple_sphere">
+  <link name="sphere">
+    <inertial>
+      <mass value="1.0" />
+      <origin xyz="0 0 0" /> 
+      <inertia  ixx="1.0" ixy="0.0"  ixz="0.0"  iyy="1.0"  iyz="0.0"  izz="1.0" />
+    </inertial>
+    <visual>
+      <!-- visual origin is defined w.r.t. link local coordinate system -->
+      <origin xyz="0 0 0" rpy="0 0 0" />
+      <geometry>
+        <sphere radius="0.05"/>
+      </geometry>
+    </visual>
+    <collision>
+      <!-- collision origin is defined w.r.t. link local coordinate system -->
+      <origin xyz="0 0 0" rpy="0 0 0" />
+      <geometry>
+        <sphere radius="0.05"/>
+      </geometry>
+    </collision>
+  </link>
+</robot>
diff --git a/scripts/test_environment.py b/scripts/test_environment.py
@@ -8,18 +8,20 @@
 sys.path.append(ROOT_DIR)
 
 
-from tprmp.envs.tasks import PalletizingBoxes # noqa
+from tprmp.envs.tasks import PickBox # noqa
 from tprmp.envs.gym import Environment  # noqa
 
 
-max_steps = 500
-env = Environment(task=PalletizingBoxes(), disp=True)
+max_steps = 1000
+env = Environment(task=PickBox(), disp=True)
+env.task.spawn_sphere(env)
 t = 0
-a = 5.
+a = 100.
 T = 500
+timeout = 20
 omega = 2 * np.pi / T
 margin = 0.0009
-speed = 0.002
+speed = 0.00005
 
 # oscillate on y-axis one T
 while t < max_steps:
@@ -28,14 +30,14 @@
     t += 1
 
 # grasp box
-box_id = env.task.goals[0][0][0]
+box_id = env.task.box_id
 target = p.getBasePositionAndOrientation(box_id)
 pos = np.array(target[0])
 pos[2] += env.task.box_size[2] / 2 + margin
 target = np.append(pos, target[1])
-env.movep(target, speed=speed)
+env.movep(target, speed=speed, timeout=timeout)
 env.ee.activate()
-env.movep(env.home_pose, speed=speed)
+env.movep(env.home_pose, speed=speed, timeout=timeout)
 env.ee.release()
 for _ in range(100):  # to simulate dropping
     p.stepSimulation()

diff --git a/scripts/test_tprmp.py b/scripts/test_tprmp.py
@@ -17,7 +17,7 @@
 from tprmp.demonstrations.frame import Frame  # noqa
 from tprmp.demonstrations.quaternion import q_convert_wxyz, q_from_euler, q_convert_xyzw  # noqa
 from tprmp.envs.gym import Environment # noqa
-from tprmp.envs.tasks import PalletizingBoxes # noqa
+from tprmp.envs.tasks import PickBox # noqa
 
 parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                  description='Example run: python test_tprmp.py')
@@ -57,12 +57,12 @@
     model.save(name=args.data)
 model.model.plot_model(demos, var_scale=var_scale)
 # test tprmp
-env = Environment(task=PalletizingBoxes(), disp=True, sampling_hz=sampling_hz)
-ee_pose = np.array(np.array(demos[0].traj[:, 0]))
+env = Environment(task=PickBox(), disp=True, sampling_hz=sampling_hz)
+ee_pose = np.array(demos[0].traj[:, 0])
 A, b = Demonstration.construct_linear_map(manifold, ee_pose)
 ee_frame = Frame(A, b, manifold=manifold)
-box_id = env.task.goals[0][0][0]
-position = np.array([0.5, -0.25, env.task.box_size[0] / 2])  # + np.random.uniform(low=-r, high=r) * np.array([1, 1, 0])
+box_id = env.task.box_id
+position = np.array([0.5, -0.25, env.task.box_size[1] / 2])  # + np.random.uniform(low=-r, high=r) * np.array([1, 1, 0])
 rotation = q_convert_xyzw(q_from_euler(np.array([np.pi/2, 0., 0.])))
 p.resetBasePositionAndOrientation(box_id, position, rotation)
 target = p.getBasePositionAndOrientation(box_id)
@@ -79,5 +79,5 @@
 for t in np.linspace(0, T, T * env.sampling_hz + 1):
     x = np.append(env.ee_pose[:3], q_convert_wxyz(env.ee_pose[3:]))
     dx = env.ee_vel
-    ddx = model.retrieve(x, dx, frames)
+    ddx = model.retrieve(x, dx)
     env.step(ddx)
diff --git a/scripts/test_tprmp_with_rmpflow.py b/scripts/test_tprmp_with_rmpflow.py
@@ -0,0 +1,100 @@
+import sys
+import argparse
+import numpy as np
+import pybullet as p
+from os.path import join, dirname, abspath
+import logging
+logging.basicConfig()
+logging.getLogger().setLevel(logging.INFO)
+
+ROOT_DIR = join(dirname(abspath(__file__)), '..')
+sys.path.append(ROOT_DIR)
+from tprmp.utils.loading import load_demos  # noqa
+from tprmp.visualization.demonstration import plot_demo  # noqa
+from tprmp.visualization.dynamics import plot_heatmap_3d  # noqa
+from tprmp.visualization.models import plot_gmm  # noqa
+from tprmp.models.tp_rmp import TPRMP  # noqa
+from tprmp.models.rmp_tree import RMPLeaf, RMPNode, RMPRoot  # noqa
+from tprmp.models.rmp_models import CollisionAvoidance  # noqa
+from tprmp.demonstrations.manifold import Manifold  # noqa
+from tprmp.demonstrations.base import Demonstration  # noqa
+from tprmp.demonstrations.frame import Frame  # noqa
+from tprmp.demonstrations.quaternion import q_convert_wxyz, q_from_euler, q_convert_xyzw  # noqa
+from tprmp.envs.gym import Environment # noqa
+from tprmp.envs.tasks import PickBox # noqa
+
+parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+                                 description='Example run: python test_tprmp.py')
+parser.add_argument('--task', help='The task folder', type=str, default='pick')
+parser.add_argument('--data', help='The data file', type=str, default='sample2.p')
+args = parser.parse_args()
+
+DATA_DIR = join(ROOT_DIR, 'data', 'tasks', args.task, 'demos')
+data_file = join(DATA_DIR, args.data)
+# parameters
+T = 100
+R = 0.06
+r = 0.01
+model = TPRMP.load(args.task, model_name=args.data)
+manifold = model.model.manifold
+dt = model.model.dt
+# environment
+env = Environment(task=PickBox(R=R), disp=True, sampling_hz=int(1/dt))
+env.task.spawn_sphere(env)
+sphere_pose, _ = p.getBasePositionAndOrientation(env.task.sphere_id)
+# build rmp tree
+root = RMPRoot('C_space', manifold=Manifold.get_euclidean_manifold(len(env.joints)))
+
+
+def tprmp_psi(q):
+    pose = np.array(env.ee_pose)
+    pose[3:] = q_convert_wxyz(pose[3:])
+    return pose
+
+
+def tprmp_J(q):
+    if isinstance(q, np.ndarray):
+        q = q.tolist()
+    zero_v = [0.] * len(q)
+    joint_states = (q, zero_v, zero_v)
+    J_pos, J_rot = env.compute_ee_jacobian(joint_states=joint_states)
+    J = np.append(np.array(J_pos), np.array(J_rot), axis=0)
+    return J
+
+
+def ws_J(q):
+    if isinstance(q, np.ndarray):
+        q = q.tolist()
+    zero_v = [0.] * len(q)
+    joint_states = (q, zero_v, zero_v)
+    J_pos, _ = env.compute_ee_jacobian(joint_states=joint_states)
+    return np.array(J_pos)
+
+
+ws_node = RMPNode('R^3_space', parent=root, manifold=Manifold.get_euclidean_manifold(3), psi=lambda x: env.ee_pose[:3], J=ws_J)
+ca_node = CollisionAvoidance('CA_space', parent=ws_node, c=np.array(sphere_pose), R=env.task.R)
+tprmp_node = RMPLeaf('TPRMP_space', model.rmp, parent=root, manifold=manifold, psi=tprmp_psi, J=tprmp_J)
+# init start & end pose
+start_pose = np.array([4.35803125e-1, 1.09041607e-1, 2.90120033e-1, 9.93708392e-1, -1.76660117e-4, 1.11998216e-1, 9.23757958e-6])
+A, b = Demonstration.construct_linear_map(manifold, start_pose)
+ee_frame = Frame(A, b, manifold=manifold)
+box_id = env.task.box_id
+position = np.array([0.5, -0.25, env.task.box_size[1] / 2])  # + np.random.uniform(low=-r, high=r) * np.array([1, 1, 0])
+rotation = q_convert_xyzw(q_from_euler(np.array([np.pi/2, 0., 0.])))
+p.resetBasePositionAndOrientation(box_id, position, rotation)
+target = p.getBasePositionAndOrientation(box_id)
+obj_pose = np.append(position, q_convert_wxyz(rotation))
+A, b = Demonstration.construct_linear_map(manifold, obj_pose)
+obj_frame = Frame(A, b, manifold=manifold)
+frames = {'ee_frame': ee_frame, 'obj_frame': obj_frame}
+plot_gmm(model.model, frames, var_scale=model.var_scale, new_fig=True, show=True)
+model.generate_global_gmm(frames)
+curr = np.array(start_pose)
+curr[3:] = q_convert_xyzw(curr[3:])
+env.setp(curr)
+env._ee_pose = curr
+env._config, _, _ = env.get_joint_states(np_array=True)
+# execution
+for t in np.linspace(0, T, T * env.sampling_hz + 1):
+    ddq = root.solve(env.config, env.config_vel)
+    env.step(ddq, return_data=False, config_space=True)
diff --git a/tprmp/envs/gym.py b/tprmp/envs/gym.py
@@ -194,6 +194,10 @@ def movep(self, pose, speed=0.01, timeout=5, direct=False, wait=0.):
     def setp(self, pose):
         """This should work with p.stepSimulation()"""
         targj = self.solve_ik(pose)
+        self.setj(targj)
+
+    def setj(self, targj):
+        """This should work with p.stepSimulation()"""
         for i in range(len(self.joints)):
             p.resetJointState(self.ur5, self.joints[i], targj[i])
         if not self.real_time_step:
@@ -274,8 +278,10 @@ def step(self, action=None, return_data=True, config_space=False):
             j += j_vel * dt
             self._config = j
             self._config_vel = j_vel
+            self.setj(j)
             ee = self.forward_kinematics()
-            joint_states = (j, j_vel, np.zeros_like(j))
+            zeros_v = [0.] * len(j)
+            joint_states = (j.tolist(), zeros_v, zeros_v)
             J_pos, J_rot = self.compute_ee_jacobian(joint_states)
             J_pos, J_rot = np.array(J_pos), np.array(J_rot)
             self._ee_pose = ee
@@ -368,9 +374,6 @@ def record_trajectory(self):
             time.sleep(1 / self.sampling_hz)
         traj = np.vstack(traj).T
         traj_vel = np.vstack(traj_vel).T
-        # record traj at self.sampling_hz
-        # traj = traj[:, ::self.sampling_hz]
-        # traj_vel = traj_vel[:, ::self.sampling_hz]
         return traj, traj_vel
 
     @property

diff --git a/tprmp/envs/tasks.py b/tprmp/envs/tasks.py
@@ -14,6 +14,7 @@
 _path_file = dirname(realpath(__file__))
 ASSETS_PATH = join(_path_file, '..', '..', 'data', 'assets')
 BOX_URDF = join(ASSETS_PATH, 'box', 'box-template.urdf')
+SPHERE_URDF = join(ASSETS_PATH, 'sphere', 'sphere-template.urdf')
 PALLET_URDF = join(ASSETS_PATH, 'pallet', 'pallet.urdf')
 
 
@@ -90,6 +91,52 @@ def get_object_points(self, obj):
         return np.vstack((xv.reshape(1, -1), yv.reshape(1, -1), zv.reshape(1, -1)))
 
 
+class PickBox(Task):
+    """Pick box task with collision avoidance."""
+    logger = logging.getLogger(__name__)
+
+    def __init__(self, box_size=None, R=0.03):
+        super(PickBox, self).__init__()
+        self.max_steps = 30
+        if box_size is None:
+            box_size = np.array([17. / 300, 0.09, 17. / 300])
+        self.box_id = None
+        self.sphere_id = None
+        self.box_size = box_size
+        self.R = R
+
+    def reset(self, env, wait=50, random=False, r=0.01):
+        super(PickBox, self).reset(env)
+        if random:
+            theta = np.random.random() * 2 * np.pi
+            pos_rand = np.random.uniform(low=-r, high=r) * np.array([1, 1, 0])
+        else:
+            theta = 0
+            pos_rand = np.zeros(3)
+        position = np.array([0.5, -0.25, self.box_size[2] / 2]) + pos_rand
+        rotation = q_convert_xyzw(q_from_euler(np.array([0., 0., theta])))
+        pose = np.append(position, rotation)
+        urdf = self.fill_template(BOX_URDF, {'DIM': self.box_size})
+        self.box_id = env.add_object(urdf, pose)
+        os.remove(urdf)
+        self.color_random_brown(self.box_id)
+        # Wait until spawned box settles.
+        for _ in range(wait):
+            p.stepSimulation()
+        return self.box_id
+
+    def spawn_sphere(self, env, sphere_pose=None, static=True):
+        if sphere_pose is None:
+            sphere_pose = np.array([0.5, -0.08, 0.13, 0., 0., 0., 1.])
+        urdf = self.fill_template(SPHERE_URDF, {'RADIUS': [self.R], 'MASS': [0. if static else 1.]})
+        self.sphere_id = env.add_object(urdf, sphere_pose)
+        os.remove(urdf)
+
+    def reward(self):
+        reward, info = super().reward()
+        return reward, info
+
+
 class PalletizingBoxes(Task):
     """Palletizing Task."""
     logger = logging.getLogger(__name__)

diff --git a/tprmp/models/rmp_models.py b/tprmp/models/rmp_models.py
@@ -38,7 +38,7 @@ def rmp_func(x, dx, M_limit=1e5, f_limit=1e10):
             f = - grad_Phi - xi - Bdx
             f = min(max(f, -f_limit), f_limit)
             return M, f
-        RMPLeaf.__init__(self, name, parent, rmp_func, manifold=Manifold.get_euclidean_manifold(1))
+        RMPLeaf.__init__(self, name, rmp_func, parent=parent, manifold=Manifold.get_euclidean_manifold(1))
         # mappings f: R^3 -> R
         self.psi = lambda x: np.array(norm(x - self.c) / self.R - 1)  # noqa
         self.J = lambda x: (1. / (self.R * norm(x - self.c))) * (x - self.c).T  # noqa

diff --git a/tprmp/models/rmp_tree.py b/tprmp/models/rmp_tree.py
@@ -40,7 +40,7 @@ def pushforward(self):
         if self.psi is not None and self.parent.x is not None:
             self.x = self.psi(self.parent.x)
         if self.J is not None and self.parent.dx is not None:
-            self.dx = self.J @ self.parent.dx
+            self.dx = np.dot(self.J(self.parent.x), self.parent.dx)
         for child in self.children:
             child.pushforward()
 
@@ -56,9 +56,10 @@ def pullback(self):
             if child.f is not None and child.M is not None:
                 f_term = child.f
                 if child.J_dot is not None:
-                    f_term -= child.M @ child.J_dot @ self.dx
-                f += child.J @ f_term
-                M += child.J.T @ child.M @ child.J
+                    f_term -= np.dot(np.dot(child.M, child.J_dot(self.x, self.dx)), self.dx)
+                J = child.J(self.x)
+                f += np.dot(J.T, f_term)
+                M += np.dot(np.dot(J.T, child.M), J)
         self.f = f
         self.M = M
 
@@ -73,17 +74,18 @@ def set_root_state(self, x, dx):
 
     def pushforward(self):
         logger.debug(f'{self.name}: Root pushforward')
-        [child.pushforward() for child in self.children]
+        for child in self.children:
+            child.pushforward()
 
     def resolve(self):
         logger.debug(f'{self.name}: Root pullback')
-        self.a = np.linalg.pinv(self.M) @ self.f
+        self.a = np.dot(np.linalg.pinv(self.M), self.f)
         return self.a
 
-    def solve(self, x, dx, **kwarg):
+    def solve(self, x, dx):
         self.set_root_state(x, dx)
         self.pushforward()
-        self.pullback(**kwarg)
+        self.pullback()
         return self.resolve()
 
 
@@ -92,9 +94,6 @@ def __init__(self, name, rmp_func, parent=None, manifold=None, psi=None, J=None,
         RMPNode.__init__(self, name, parent, manifold, psi, J, J_dot)
         self.rmp_func = rmp_func
 
-    def rmp(self, **kwarg):
-        self.M, self.f = self.rmp_func(self.x, self.dx, **kwarg)
-
-    def pullback(self, **kwarg):
+    def pullback(self):
         logger.debug(f'{self.name}: leaf pullback')
-        self.rmp(**kwarg)
+        self.M, self.f = self.rmp_func(self.x, self.dx)