mirror of
https://github.com/huggingface/lerobot.git
synced 2026-06-03 12:21:27 +00:00
Modified kinematics code to be independant of drive mode
Modified gym_manipulator.py and find_joint_limits to adhere to the refactor of robot devices Modified the configuration of envs to take into account the refactor
This commit is contained in:
committed by
AdilZouitine
parent
ba477e81ce
commit
2475645f5f
@@ -14,122 +14,93 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import argparse
|
||||
"""
|
||||
Simple script to control a robot from teleoperation.
|
||||
|
||||
Example:
|
||||
|
||||
```shell
|
||||
python -m lerobot.scripts.server.find_joint_limits \
|
||||
--robot.type=so100_follower \
|
||||
--robot.port=/dev/tty.usbmodem58760431541 \
|
||||
--robot.id=black \
|
||||
--teleop.type=so100_leader \
|
||||
--teleop.port=/dev/tty.usbmodem58760431551 \
|
||||
--teleop.id=blue
|
||||
```
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import draccus
|
||||
|
||||
from lerobot.common.robot_devices.control_utils import is_headless
|
||||
from lerobot.common.robot_devices.robots.configs import RobotConfig
|
||||
from lerobot.common.robot_devices.robots.utils import make_robot_from_config
|
||||
from lerobot.configs import parser
|
||||
from lerobot.scripts.server.kinematics import RobotKinematics
|
||||
from lerobot.common.robots import ( # noqa: F401
|
||||
RobotConfig,
|
||||
koch_follower,
|
||||
make_robot_from_config,
|
||||
so100_follower,
|
||||
so100_follower_end_effector,
|
||||
)
|
||||
from lerobot.common.teleoperators import (
|
||||
TeleoperatorConfig,
|
||||
make_teleoperator_from_config,
|
||||
)
|
||||
|
||||
from dataclasses import dataclass
|
||||
from lerobot.common.teleoperators import gamepad, koch_leader, so100_leader # noqa: F401
|
||||
from lerobot.common.model.kinematics import RobotKinematics
|
||||
|
||||
|
||||
def find_joint_bounds(
|
||||
robot,
|
||||
control_time_s=30,
|
||||
display_cameras=False,
|
||||
):
|
||||
if not robot.is_connected:
|
||||
robot.connect()
|
||||
@dataclass
|
||||
class FindJointLimitsConfig:
|
||||
teleop: TeleoperatorConfig
|
||||
robot: RobotConfig
|
||||
# Limit the maximum frames per second. By default, no limit.
|
||||
fps: int | None = None
|
||||
teleop_time_s: float | None = None
|
||||
# Display all cameras on screen
|
||||
display_data: bool = False
|
||||
|
||||
start_episode_t = time.perf_counter()
|
||||
pos_list = []
|
||||
while True:
|
||||
observation, action = robot.teleop_step(record_data=True)
|
||||
@draccus.wrap()
|
||||
def find_joint_and_ee_bounds(cfg: FindJointLimitsConfig):
|
||||
teleop = make_teleoperator_from_config(cfg.teleop)
|
||||
robot = make_robot_from_config(cfg.robot)
|
||||
|
||||
# Wait for 5 seconds to stabilize the robot initial position
|
||||
if time.perf_counter() - start_episode_t < 5:
|
||||
continue
|
||||
|
||||
pos_list.append(robot.follower_arms["main"].read("Present_Position"))
|
||||
|
||||
if display_cameras and not is_headless():
|
||||
image_keys = [key for key in observation if "image" in key]
|
||||
for key in image_keys:
|
||||
cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
|
||||
cv2.waitKey(1)
|
||||
|
||||
if time.perf_counter() - start_episode_t > control_time_s:
|
||||
max = np.max(np.stack(pos_list), 0)
|
||||
min = np.min(np.stack(pos_list), 0)
|
||||
print(f"Max angle position per joint {max}")
|
||||
print(f"Min angle position per joint {min}")
|
||||
break
|
||||
|
||||
|
||||
def find_ee_bounds(
|
||||
robot,
|
||||
control_time_s=30,
|
||||
display_cameras=False,
|
||||
):
|
||||
if not robot.is_connected:
|
||||
robot.connect()
|
||||
teleop.connect()
|
||||
robot.connect()
|
||||
|
||||
start_episode_t = time.perf_counter()
|
||||
ee_list = []
|
||||
pos_list = []
|
||||
robot_type = cfg.robot.type.split("_")[0]
|
||||
kinematics = RobotKinematics(robot_type)
|
||||
control_time_s = 30
|
||||
while True:
|
||||
observation, action = robot.teleop_step(record_data=True)
|
||||
action = teleop.get_action()
|
||||
robot.send_action(action)
|
||||
observation = robot.get_observation()
|
||||
|
||||
# Wait for 5 seconds to stabilize the robot initial position
|
||||
if time.perf_counter() - start_episode_t < 5:
|
||||
continue
|
||||
|
||||
kinematics = RobotKinematics(robot.robot_type)
|
||||
joint_positions = robot.follower_arms["main"].read("Present_Position")
|
||||
print(f"Joint positions: {joint_positions}")
|
||||
joint_positions = robot.bus.sync_read("Present_Position")
|
||||
joint_positions = np.array([joint_positions[key] for key in joint_positions.keys()])
|
||||
ee_list.append(kinematics.fk_gripper_tip(joint_positions)[:3, 3])
|
||||
|
||||
if display_cameras and not is_headless():
|
||||
image_keys = [key for key in observation if "image" in key]
|
||||
for key in image_keys:
|
||||
cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
|
||||
cv2.waitKey(1)
|
||||
pos_list.append(joint_positions)
|
||||
|
||||
if time.perf_counter() - start_episode_t > control_time_s:
|
||||
max = np.max(np.stack(ee_list), 0)
|
||||
min = np.min(np.stack(ee_list), 0)
|
||||
print(f"Max ee position {max}")
|
||||
print(f"Min ee position {min}")
|
||||
max_ee = np.max(np.stack(ee_list), 0)
|
||||
min_ee = np.min(np.stack(ee_list), 0)
|
||||
max_pos = np.max(np.stack(pos_list), 0)
|
||||
min_pos = np.min(np.stack(pos_list), 0)
|
||||
print(f"Max ee position {max_ee}")
|
||||
print(f"Min ee position {min_ee}")
|
||||
print(f"Max joint pos position {max_pos}")
|
||||
print(f"Min joint pos position {min_pos}")
|
||||
break
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Create argparse for script-specific arguments
|
||||
parser = argparse.ArgumentParser(add_help=False) # Set add_help=False to avoid conflict
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
type=str,
|
||||
default="joint",
|
||||
choices=["joint", "ee"],
|
||||
help="Mode to run the script in. Can be 'joint' or 'ee'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--control-time-s",
|
||||
type=int,
|
||||
default=30,
|
||||
help="Time step to use for control.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--robot-type",
|
||||
type=str,
|
||||
default="so100",
|
||||
help="Robot type (so100, koch, aloha, etc.)",
|
||||
)
|
||||
|
||||
# Only parse known args, leaving robot config args for Hydra if used
|
||||
args = parser.parse_args()
|
||||
|
||||
# Create robot with the appropriate config
|
||||
robot_config = RobotConfig.get_choice_class(args.robot_type)(mock=False)
|
||||
robot = make_robot_from_config(robot_config)
|
||||
|
||||
if args.mode == "joint":
|
||||
find_joint_bounds(robot, args.control_time_s)
|
||||
elif args.mode == "ee":
|
||||
find_ee_bounds(robot, args.control_time_s)
|
||||
|
||||
if robot.is_connected:
|
||||
robot.disconnect()
|
||||
find_joint_and_ee_bounds()
|
||||
|
||||
Reference in New Issue
Block a user