gello_software/gello/robots/dynamixel.py

134 lines
4.7 KiB
Python

from typing import Dict, Optional, Sequence, Tuple
import numpy as np
from gello.robots.robot import Robot
class DynamixelRobot(Robot):
"""A class representing a UR robot."""
def __init__(
self,
joint_ids: Sequence[int],
joint_offsets: Optional[Sequence[float]] = None,
joint_signs: Optional[Sequence[int]] = None,
real: bool = False,
port: str = "/dev/ttyUSB0",
baudrate: int = 57600,
gripper_config: Optional[Tuple[int, float, float]] = None,
start_joints: Optional[np.ndarray] = None,
):
from gello.dynamixel.driver import (
DynamixelDriver,
DynamixelDriverProtocol,
FakeDynamixelDriver,
)
print(f"attempting to connect to port: {port}")
self.gripper_open_close: Optional[Tuple[float, float]]
if gripper_config is not None:
assert joint_offsets is not None
assert joint_signs is not None
# joint_ids.append(gripper_config[0])
# joint_offsets.append(0.0)
# joint_signs.append(1)
joint_ids = tuple(joint_ids) + (gripper_config[0],)
joint_offsets = tuple(joint_offsets) + (0.0,)
joint_signs = tuple(joint_signs) + (1,)
self.gripper_open_close = (
gripper_config[1] * np.pi / 180,
gripper_config[2] * np.pi / 180,
)
else:
self.gripper_open_close = None
self._joint_ids = joint_ids
self._driver: DynamixelDriverProtocol
if joint_offsets is None:
self._joint_offsets = np.zeros(len(joint_ids))
else:
self._joint_offsets = np.array(joint_offsets)
if joint_signs is None:
self._joint_signs = np.ones(len(joint_ids))
else:
self._joint_signs = np.array(joint_signs)
assert len(self._joint_ids) == len(self._joint_offsets), (
f"joint_ids: {len(self._joint_ids)}, "
f"joint_offsets: {len(self._joint_offsets)}"
)
assert len(self._joint_ids) == len(self._joint_signs), (
f"joint_ids: {len(self._joint_ids)}, "
f"joint_signs: {len(self._joint_signs)}"
)
assert np.all(
np.abs(self._joint_signs) == 1
), f"joint_signs: {self._joint_signs}"
if real:
self._driver = DynamixelDriver(joint_ids, port=port, baudrate=baudrate)
self._driver.set_torque_mode(False)
else:
self._driver = FakeDynamixelDriver(joint_ids)
self._torque_on = False
self._last_pos = None
self._alpha = 0.99
if start_joints is not None:
# loop through all joints and add +- 2pi to the joint offsets to get the closest to start joints
new_joint_offsets = []
current_joints = self.get_joint_state()
assert current_joints.shape == start_joints.shape
if gripper_config is not None:
current_joints = current_joints[:-1]
start_joints = start_joints[:-1]
for c_joint, s_joint, joint_offset in zip(
current_joints, start_joints, self._joint_offsets
):
new_joint_offsets.append(
np.pi * 2 * np.round((s_joint - c_joint) / (2 * np.pi))
+ joint_offset
)
if gripper_config is not None:
new_joint_offsets.append(self._joint_offsets[-1])
self._joint_offsets = np.array(new_joint_offsets)
def num_dofs(self) -> int:
return len(self._joint_ids)
def get_joint_state(self) -> np.ndarray:
pos = (self._driver.get_joints() - self._joint_offsets) * self._joint_signs
assert len(pos) == self.num_dofs()
if self.gripper_open_close is not None:
# map pos to [0, 1]
g_pos = (pos[-1] - self.gripper_open_close[0]) / (
self.gripper_open_close[1] - self.gripper_open_close[0]
)
g_pos = min(max(0, g_pos), 1)
pos[-1] = g_pos
if self._last_pos is None:
self._last_pos = pos
else:
# exponential smoothing
pos = self._last_pos * (1 - self._alpha) + pos * self._alpha
self._last_pos = pos
return pos
def command_joint_state(self, joint_state: np.ndarray) -> None:
self._driver.set_joints((joint_state + self._joint_offsets).tolist())
def set_torque_mode(self, mode: bool):
if mode == self._torque_on:
return
self._driver.set_torque_mode(mode)
self._torque_on = mode
def get_observations(self) -> Dict[str, np.ndarray]:
return {"joint_state": self.get_joint_state()}