robotics-visual-lab / robots.py
askuric's picture
askuric HF Staff
update to load talos first
7a6243f
import numpy as np
import example_robot_data as erd
import pinocchio as pin
import meshcat
import meshcat.geometry as g
from pinocchio.visualize import MeshcatVisualizer
import time
def _actuated_dof_indices(model):
idxs, labels, bounds = [], [], []
for j in range(1, len(model.joints)):
J = model.joints[j]
if J.shortname() == "JointModelFreeFlyer":
continue
if J.nq == 1:
i0 = J.idx_q
lo = float(model.lowerPositionLimit[i0])
hi = float(model.upperPositionLimit[i0])
if not np.isfinite(lo): lo = -3.14
if not np.isfinite(hi): hi = 3.14
idxs.append(i0)
labels.append(model.names[j])
bounds.append((lo, hi))
return idxs, labels, bounds
class LiveRobot:
"""
Starts a local MeshCat ZMQ+Web server and keeps a persistent connection.
The web UI is proxied at /meshcat/ by nginx (see nginx.conf).
"""
def __init__(self, robot_name="talos"):
self.vis = meshcat.Visualizer(zmq_url="tcp://127.0.0.1:6001")
print(f"Zmq URL: {self.vis.window.zmq_url}, web URL: {self.vis.window.web_url}")
self.rdata = erd.load(robot_name)
self.model = self.rdata.model
self.data = self.model.createData()
self.vis.delete()
time.sleep(0.5) # wait for the robot to appear
self.viz = MeshcatVisualizer(
self.model, self.rdata.collision_model, self.rdata.visual_model
)
self.viz.initViewer(self.vis)
self.viz.clean()
time.sleep(0.5) # wait for the robot to appear
self.viz.loadViewerModel(rootNodeName="robot")
self.q = pin.neutral(self.model)
try:
self.viz.display(self.q)
except Exception as e:
print(f"Error setting robot joints: {e}")
self.idxs, self.labels, self.bounds = _actuated_dof_indices(self.model)
def set_joints(self, vals):
q = self.q.copy()
for v, qi in zip(vals, self.idxs):
q[qi] = float(v)
try:
self.viz.display(q)
self.q = q
except Exception as e:
print(f"Error setting robot joints: {e}")
def neutral(self):
self.q = pin.neutral(self.model)
self.viz.display(self.q)
return [float(self.q[i]) for i in self.idxs]
def iframe(self, width="100%", height=640):
# nginx proxies the MeshCat web app at /meshcat/
return f'<iframe src="/static/" style="width:{width};height:{height}px;border:0"></iframe>'