reachy-mini-viewer / robot.py
askuric's picture
askuric HF Staff
update speed
ca40f21
import numpy as np
import time
import meshcat
import meshcat.geometry as g
import logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
import pinocchio as pin
from pinocchio.visualize import MeshcatVisualizer
from placo_utils.tf import tf
from placo_utils.visualization import frame_viz, robot_frame_viz, robot_viz
import reachy_mini.utils.constants
from reachy_mini.kinematics.placo_kinematics import PlacoKinematics
urdf_path = reachy_mini.utils.constants.URDF_ROOT_PATH
robot = pin.RobotWrapper.BuildFromURDF("./robot_only_outside.urdf", urdf_path, pin.JointModelFreeFlyer())
solver = PlacoKinematics(urdf_path, 0.02)
class LiveReachyMini:
"""
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):
self.vis = meshcat.Visualizer(zmq_url="tcp://127.0.0.1:6001")
logging.info(f"Zmq URL: {self.vis.window.zmq_url}, web URL: {self.vis.window.web_url}")
self.rdata = robot
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
)
# preare the visualisation
self.viz.initViewer(self.vis)
# make it pretty
self.vis["/Background"].set_property("top_color", [1, 1, 1])
self.vis["/Background"].set_property("bottom_color", [1, 1, 1])
self.vis["/Axes"].set_property("visible", False)
self.vis['/Grid'].set_property("visible", False)
self.vis["/Cameras/default"].set_transform(
tf.translation_matrix([0, 0, 0.3])
)
self.vis["/Cameras/default/rotated/<object>"].set_property("zoom", 5.0)
#self.vis["/Cameras/default/rotated/<object>"].set_property("position", [1, 0.5, 0])
self.viz.clean()
time.sleep(0.5) # wait for the robot to appear
self.viz.loadViewerModel(rootNodeName="robot")
solver.ik(np.eye(4), body_yaw=0, no_iterations=16)
self.q = solver.robot.state.q.copy()
try:
self.viz.display(self.q)
except Exception as e:
logging.error(f"Error setting robot joints: {e}")
# setting the target
# head roll, pitch, yaw,
# head x,y,z,
# body yaw,
# antennas left/right
def set_target(self, vals):
q = self.q.copy()
# solve roll, pitch, yaw for head
r, p, yaw = vals[1:4]
x, y, z = vals[4:7]
T_head = tf.translation_matrix((x, y, z)) @ tf.euler_matrix(
r, p, yaw
)
#T_head[2,3] += t # add some vertical bobbing
solver.ik(T_head, body_yaw=vals[0], no_iterations=5)
q = solver.robot_ik.state.q.copy()
# set antennas
q [ self.model.getJointId("left_antenna") + 6 -1 ] = vals[7]
q [ self.model.getJointId("right_antenna") + 6 -1 ] = vals[8]
try:
s = time.time()
self.viz.display(q)
self.q = q
logging.debug("elapsed time after display: {:.2f} ms".format((time.time() - s)*1000))
except Exception as e:
logging.error(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>'