Skip to content

Commit

Permalink
Merge branch 'topic/ci-conda-python-stubs' of github.com:cmastalli/cr…
Browse files Browse the repository at this point in the history 
…ocoddyl into topic/ci-conda-python-stubs
  • Loading branch information
@cmastalli
cmastalli committed Mar 14, 2024
2 parents 0b442e5 + 98e9bf6 commit b81d4ce
Show file tree
Hide file tree
Showing 2 changed files with 112 additions and 41 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
## [Unreleased]

* Tested Python stubs in Conda CI in https://github.com/loco-3d/crocoddyl/pull/1228
* 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
Expand Down
152 changes: 111 additions & 41 deletions bindings/python/crocoddyl/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand All @@ -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()))
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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"):
Expand All @@ -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,
Expand All @@ -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

Expand Down

0 comments on commit b81d4ce

Please sign in to comment.