89 lines
2.8 KiB
Python
Raw Permalink Normal View History

2025-07-28 12:35:44 +08:00
import argparse
import time
import warnings
import numpy as np
import placo
from placo_utils.visualization import robot_viz
from placo_utils.tf import tf
import time
import pickle
import FramesViewer.utils as fv_utils
warnings.filterwarnings("ignore")
model_filename = "../../mini_bdx/robots/bdx/robot.urdf"
robot = placo.RobotWrapper(model_filename)
robot.set_joint_limits("left_knee", -2, -0.01)
robot.set_joint_limits("right_knee", -2, -0.01)
solver = placo.KinematicsSolver(robot)
# solver.mask_fbase(True)
left_foot_task = solver.add_frame_task("left_foot_frame", np.eye(4))
left_foot_task.configure("left_foot_frame", "soft", 1.0)
right_foot_task = solver.add_frame_task("right_foot_frame", np.eye(4))
right_foot_task.configure("right_foot_frame", "soft", 1.0)
T_world_trunk = np.eye(4)
T_world_trunk = fv_utils.rotateInSelf(T_world_trunk, [0, -2, 0])
trunk_task = solver.add_frame_task("trunk", T_world_trunk)
trunk_task.configure("trunk", "hard")
T_world_head = np.eye(4)
T_world_head[:3, 3][2] = 0.1
head_task = solver.add_frame_task(
"head",
T_world_head,
)
head_task.configure("head", "soft")
robot.update_kinematics()
move = []
viz = robot_viz(robot)
while True: # some main loop
# Update tasks data here
# left_foot_task.T_world_frame = tf.translation_matrix(
# [0, -0.12 / 2, np.sin(time.time())]
# )
z_offset = 0.015 * np.sin(2 * time.time())
left_foot_task.T_world_frame = tf.translation_matrix(
[-0.005, 0.12 / 2, -0.17 + z_offset]
)
right_foot_task.T_world_frame = tf.translation_matrix(
[-0.005, -0.12 / 2, -0.17 + z_offset]
)
head_task.T_world_frame = tf.translation_matrix(
[0.05 * np.sin(2 * np.pi * 2 * time.time()), 0, 0.1]
)
# Solve the IK
solver.solve(True)
# Update frames and jacobians
robot.update_kinematics()
angles = {
"right_hip_yaw": robot.get_joint("right_hip_yaw"),
"right_hip_roll": robot.get_joint("right_hip_roll"),
"right_hip_pitch": robot.get_joint("right_hip_pitch"),
"right_knee": robot.get_joint("right_knee"),
"right_ankle": robot.get_joint("right_ankle"),
"left_hip_yaw": robot.get_joint("left_hip_yaw"),
"left_hip_roll": robot.get_joint("left_hip_roll"),
"left_hip_pitch": robot.get_joint("left_hip_pitch"),
"left_knee": robot.get_joint("left_knee"),
"left_ankle": robot.get_joint("left_ankle"),
"neck_pitch": robot.get_joint("neck_pitch"),
"head_pitch": robot.get_joint("head_pitch"),
"head_yaw": robot.get_joint("head_yaw"),
"left_antenna": robot.get_joint("left_antenna"),
"right_antenna": robot.get_joint("right_antenna"),
}
move.append(angles)
time.sleep(1 / 60)
pickle.dump(move, open("move.pkl", "wb"))
# Optionally: dump the solver status
solver.dump_status()
viz.display(robot.state.q)