Fixed Rviz display

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
 
+* Fixed Rviz display in https://github.com/loco-3d/crocoddyl/pull/1227
 * Improved CI, updated cmake and fixed launch file in https://github.com/loco-3d/crocoddyl/pull/1220
 * Introduced a Rviz display in https://github.com/loco-3d/crocoddyl/pull/1216
 * Enabled display of thrust and simplied displayers code in https://github.com/loco-3d/crocoddyl/pull/1215

bindings/python/crocoddyl/__init__.py

@@ -29,6 +29,7 @@ def rotationMatrixFromTwoVectors(a, b):
 class DisplayAbstract(ABC):
     def __init__(self, robot, rate=-1, freq=1):
         self.robot = robot
+        self._nv_root = 0.0
         self.rate = rate
         self.freq = freq
         # Visual properties
@@ -57,12 +58,13 @@ def __init__(self, robot, rate=-1, freq=1):
     def displayFromSolver(self, solver, factor=1.0):
         if not self._init:
             self.init(solver)
+        dts = self.getDurationSequenceFromSolver(solver)
         fs = self.getForceTrajectoryFromSolver(solver)
         ps = self.getFrameTrajectoryFromSolver(solver)
         rs = self.getThrustTrajectoryFromSolver(solver)
         models = [*solver.problem.runningModels.tolist(), solver.problem.terminalModel]
         dts = [m.dt if hasattr(m, "differential") else 0.0 for m in models]
-        self.display(solver.xs, [], rs, ps, [], fs, [], dts, factor)
+        self.display(solver.xs, dts, rs, ps, [], fs, [], dts, factor)
 
     def display(
         self, xs, us=[], rs=[], ps=[], pds=[], fs=[], ss=[], dts=[], factor=1.0
@@ -165,6 +167,32 @@ def displayThrustForce(self, r):
         """Display the thrust force"""
         raise NotImplementedError("Not implemented yet.")
 
+    def getTimeSequenceFromSolver(self, solver):
+        dts = self.getDurationSequenceFromSolver(solver)
+        return self._getTimeSequence(dts)
+
+    def getDurationSequenceFromSolver(self, solver):
+        return [
+            m.dt if hasattr(m, "differential") else 0.0
+            for m in solver.problem.runningModels
+        ]
+
+    def getJointTorquesTrajectoryFromSolver(self, solver):
+        us = []
+        for i in range(solver.problem.T):
+            data = solver.problem.runningDatas[i]
+            if hasattr(data, "differential"):
+                us.append(data.differential.multibody.actuation.tau[self._nv_root :])
+            elif isinstance(data, ActionDataImpulseFwdDynamics):
+                us.append(
+                    np.zeros(solver.problem.runningModels[i].state.nv - self._nv_root)
+                )
+            else:
+                us.append(solver.us[i][self._nv_root :])
+        nu = solver.problem.runningModels[i].state.nv - self._nv_root
+        us.append(np.zeros(nu + 1)[1:])  # TODO: temporal patch to fix bug in pybind11
+        return us
+
     def getForceTrajectoryFromSolver(self, solver):
         if len(self.frameTrajNames) == 0:
             return None
@@ -249,6 +277,13 @@ def getFrameTrajectoryFromSolver(self, solver):
                     p.append(np.asarray(pose.translation.T).reshape(-1).tolist())
         return ps
 
+    def _getTimeSequence(self, dts):
+        t, ts = 0.0, [0.0]
+        for dt in dts:
+            t += dt
+            ts.append(t)
+        return ts
+
     def _hasContacts(self, data):
         if hasattr(data, "differential"):
             if isinstance(
@@ -668,6 +703,10 @@ def __init__(
         freq=1,
     ):
         DisplayAbstract.__init__(self, robot, rate, freq)
+        # Disable garbage collection for better timings.
+        import gc
+
+        gc.disable()
         # Import ROS modules
         self.ROS_VERSION = int(os.environ["ROS_VERSION"])
         if self.ROS_VERSION == 2:
@@ -691,7 +730,7 @@ def __init__(
             rospy.set_param("use_sim_time", True)
             rospy.set_param(
                 "robot_description", URDF.from_xml_file(robot.urdf).to_xml_string()
-            )
+            )  # TODO: hard code robot.urdf because we cannot convert a Pinocchio model into an URDF
             filename = os.path.join(
                 os.path.dirname(
                     os.path.abspath(inspect.getfile(inspect.currentframe()))
@@ -725,48 +764,67 @@ def __init__(
             robot.model, "whole_body_trajectory", "map"
         )
         root_id = crocoddyl_ros.getRootJointId(self.robot.model)
-        self._nv_root = self.robot.model.joints[root_id].nv
+        self._nv_root = (
+            self.robot.model.joints[root_id].nv
+            if self.robot.model.frames[root_id].name != "universe"
+            else 0
+        )
 
     def __del__(self):
         if self.ROS_VERSION == 1:
             self.roscore.shutdown()
             self.rviz.shutdown()
 
     def displayFromSolver(self, solver, factor=1.0):
+        if not self._init:
+            self.init(solver)
         xs = solver.xs
-        us = []
-        for i in range(solver.problem.T):
-            us.append(solver.us[i][self._nv_root :])
-        us.append(np.zeros(0))
+        us = self.getJointTorquesTrajectoryFromSolver(solver)
+        dts = self.getDurationSequenceFromSolver(solver)
         models = [*solver.problem.runningModels.tolist(), solver.problem.terminalModel]
         dts = [m.dt if hasattr(m, "differential") else 0.0 for m in models]
         ps, pds = self.getFrameTrajectoryFromSolver(solver)
         fs, ss = self.getForceTrajectoryFromSolver(solver)
-        self.display(xs, us, [], ps, pds, fs, ss, dts, factor)
+        self.display(xs, us, dts, ps, pds, fs, ss, dts, factor)
 
     def display(
         self, xs, us=[], rs=[], ps=[], pds=[], fs=[], ss=[], dts=[], factor=1.0
     ):
         nq = self.robot.model.nq
-        ts = [i * dt for i, dt in enumerate(dts)]
+        if not dts:
+            dts = [0.01] * len(xs)
+        ts = self._getTimeSequence(dts)[:-1]
+        if len(ps) != 0:
+            self._wtPublisher.publish(ts, xs, us, ps, pds, fs, ss)
+        else:
+            self._wtPublisher.publish(ts, xs, us)
         for i in range(len(xs)):
-            if i < len(xs) - 1:
-                t, x, u = ts[i], xs[i], us[i]
-            else:
-                t, x, u = ts[i], xs[i], np.zeros_like(us[0])
+            t, x, u = ts[i], xs[i], us[i]
             if len(ps) != 0:
                 p, pd = ps[i], pds[i]
                 f, s = fs[i], ss[i]
                 self._wsPublisher.publish(t, x[:nq], x[nq:], u, p, pd, f, s)
             else:
                 self._wsPublisher.publish(t, x[:nq], x[nq:], u)
             time.sleep(dts[i] * factor)
-        self._wtPublisher.publish(ts, xs, us, ps, pds, fs, ss)
+
+    def getFrameTrajectoryFromSolver(self, solver):
+        ps, pds = [], []
+        datas = [*solver.problem.runningDatas.tolist(), solver.problem.runningDatas[-1]]
+        for i, data in enumerate(datas):
+            if self._hasContacts(data):
+                pinocchio_data = self._getPinocchioData(data)
+                ps.append(self._get_pc(pinocchio_data))
+                pds.append(self._get_pdc(pinocchio_data))
+        return ps, pds
 
     def getForceTrajectoryFromSolver(self, solver):
         fs, ss = [], []
-        models = [*solver.problem.runningModels.tolist(), solver.problem.terminalModel]
-        datas = [*solver.problem.runningDatas.tolist(), solver.problem.terminalData]
+        models = [
+            *solver.problem.runningModels.tolist(),
+            solver.problem.runningModels[-1],
+        ]
+        datas = [*solver.problem.runningDatas.tolist(), solver.problem.runningDatas[-1]]
         for i, data in enumerate(datas):
             model = models[i]
             if self._hasContacts(data):
@@ -799,19 +857,6 @@ def _addFrameCurves(self):
     def _addFrictionCones(self):
         pass
 
-    def getFrameTrajectoryFromSolver(self, solver):
-        ps, pds = [], []
-        models = [*solver.problem.runningModels.tolist(), solver.problem.terminalModel]
-        datas = [*solver.problem.runningDatas.tolist(), solver.problem.terminalData]
-        for i, data in enumerate(datas):
-            model = models[i]
-            if self._hasContacts(data):
-                pinocchio_data = self._getPinocchioData(data)
-                _, contact_data = self._getContactModelAndData(model, data)
-                ps.append(self._get_pc(pinocchio_data, contact_data))
-                pds.append(self._get_dpc(pinocchio_data, contact_data))
-        return ps, pds
-
     def _getPinocchioData(self, data):
         if hasattr(data, "differential"):
             if hasattr(data.differential, "multibody"):
@@ -822,19 +867,24 @@ def _getPinocchioData(self, data):
         elif isinstance(data, ActionDataImpulseFwdDynamics):
             return data.multibody.pinocchio
 
-    def get_pc(self, pinocchio_data, contact_data):
+    def _get_pc(self, pinocchio_data):
+        if len(self.frameTrajNames) == 0:
+            return None
         popt = dict()
-        for contact in contact_data:
-            name = contact.key()
-            frame_id = self._state.pinocchio.getFrameId(name)
+        for frame in self.frameTrajNames:
+            frame_id = int(frame)
+            name = self.robot.model.frames[frame_id].name
+            pinocchio.updateFramePlacement(self.robot.model, pinocchio_data, frame_id)
             popt[name] = pinocchio_data.oMf[frame_id]
         return popt
 
-    def get_pdc(self, pinocchio_data, contact_data):
+    def _get_pdc(self, pinocchio_data):
+        if len(self.frameTrajNames) == 0:
+            return None
         pdopt = dict()
-        for contact in contact_data:
-            name = contact.key()
-            frame_id = self._state.pinocchio.getFrameId(name)
+        for frame in self.frameTrajNames:
+            frame_id = int(frame)
+            name = self.robot.model.frames[frame_id].name
             v = pinocchio.getFrameVelocity(
                 self.robot.model,
                 pinocchio_data,
@@ -845,16 +895,36 @@ def get_pdc(self, pinocchio_data, contact_data):
         return pdopt
 
     def _get_fc(self, contact_data):
+        import crocoddyl_ros
+
+        if len(self.frameTrajNames) == 0:
+            return None
         fc = dict()
-        for force in contact_data.todict().items():
-            name, f = force[0], force[1].f
-            fc[name] = [f, 0, 2]
+        for frame in self.frameTrajNames:
+            frame_id = int(frame)
+            name = self.robot.model.frames[frame_id].name
+            fc[name] = [
+                pinocchio.Force.Zero(),
+                crocoddyl_ros.ContactType.LOCOMOTION,
+                crocoddyl_ros.ContactStatus.SEPARATION,
+            ]
+        for contact in contact_data:
+            force = contact.data().f
+            frame_id = contact.data().frame
+            name = self.robot.model.frames[frame_id].name
+            fc[name] = [
+                force,
+                crocoddyl_ros.ContactType.LOCOMOTION,
+                crocoddyl_ros.ContactStatus.STICKING,
+            ]
         return fc
 
     def _get_sc(self, contact_data):
         sc = dict()
         self.mu = 0.7
-        for name in contact_data.todict().keys():
+        for frame in self.frameTrajNames:
+            frame_id = int(frame)
+            name = self.robot.model.frames[frame_id].name
             sc[name] = [np.array([0.0, 0.0, 1.0]), self.mu]
         return sc