initial commit, add gello software code and instructions

scripts/arm_blocks_play.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+
+import numpy as np
+import tyro
+from dm_control import composer, viewer
+
+from gello.agents.gello_agent import DynamixelRobotConfig
+from gello.dm_control_tasks.arms.ur5e import UR5e
+from gello.dm_control_tasks.manipulation.arenas.floors import Floor
+from gello.dm_control_tasks.manipulation.tasks.block_play import BlockPlay
+
+
+@dataclass
+class Args:
+    use_gello: bool = False
+
+
+config = DynamixelRobotConfig(
+    joint_ids=(1, 2, 3, 4, 5, 6),
+    joint_offsets=(
+        -np.pi / 2,
+        1 * np.pi / 2 + np.pi,
+        np.pi / 2 + 0 * np.pi,
+        0 * np.pi + np.pi / 2,
+        np.pi - 2 * np.pi / 2,
+        -1 * np.pi / 2 + 2 * np.pi,
+    ),
+    joint_signs=(1, 1, -1, 1, 1, 1),
+    gripper_config=(7, 20, -22),
+)
+
+
+def main(args: Args) -> None:
+    reset_joints_left = np.deg2rad([90, -90, -90, -90, 90, 0, 0])
+    robot = UR5e()
+    task = BlockPlay(robot, Floor(), reset_joints=reset_joints_left[:-1])
+    # task = BlockPlay(robot, Floor())
+    env = composer.Environment(task=task)
+
+    action_space = env.action_spec()
+    if args.use_gello:
+        gello = config.make_robot(
+            port="/dev/cu.usbserial-FT7WBEIA", start_joints=reset_joints_left
+        )
+
+    def policy(timestep) -> np.ndarray:
+        if args.use_gello:
+            joint_command = gello.get_joint_state()
+            joint_command = np.array(joint_command).copy()
+
+            joint_command[-1] = joint_command[-1] * 255
+            return joint_command
+        return np.random.uniform(action_space.minimum, action_space.maximum)
+
+    viewer.launch(env, policy=policy)
+
+
+if __name__ == "__main__":
+    main(tyro.cli(Args))

scripts/gello_get_offset.py

@@ -0,0 +1,98 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Tuple
+
+import numpy as np
+import tyro
+
+from gello.dynamixel.driver import DynamixelDriver
+
+MENAGERIE_ROOT: Path = Path(__file__).parent / "third_party" / "mujoco_menagerie"
+
+
+@dataclass
+class Args:
+    port: str = "/dev/ttyUSB0"
+    """The port that GELLO is connected to."""
+
+    start_joints: Tuple[float, ...] = (0, 0, 0, 0, 0, 0)
+    """The joint angles that the GELLO is placed in at (in radians)."""
+
+    joint_signs: Tuple[float, ...] = (1, 1, -1, 1, 1, 1)
+    """The joint angles that the GELLO is placed in at (in radians)."""
+
+    gripper: bool = True
+    """Whether or not the gripper is attached."""
+
+    def __post_init__(self):
+        assert len(self.joint_signs) == len(self.start_joints)
+        for idx, j in enumerate(self.joint_signs):
+            assert (
+                j == -1 or j == 1
+            ), f"Joint idx: {idx} should be -1 or 1, but got {j}."
+
+    @property
+    def num_robot_joints(self) -> int:
+        return len(self.start_joints)
+
+    @property
+    def num_joints(self) -> int:
+        extra_joints = 1 if self.gripper else 0
+        return self.num_robot_joints + extra_joints
+
+
+def get_config(args: Args) -> None:
+    joint_ids = list(range(1, args.num_joints + 1))
+    driver = DynamixelDriver(joint_ids, port=args.port, baudrate=57600)
+
+    # assume that the joint state shouold be args.start_joints
+    # find the offset, which is a multiple of np.pi/2 that minimizes the error between the current joint state and args.start_joints
+    # this is done by brute force, we seach in a range of +/- 8pi
+
+    def get_error(offset: float, index: int, joint_state: np.ndarray) -> float:
+        joint_sign_i = args.joint_signs[index]
+        joint_i = joint_sign_i * (joint_state[index] - offset)
+        start_i = args.start_joints[index]
+        return np.abs(joint_i - start_i)
+
+    for _ in range(10):
+        driver.get_joints()  # warmup
+
+    for _ in range(1):
+        best_offsets = []
+        curr_joints = driver.get_joints()
+        for i in range(args.num_robot_joints):
+            best_offset = 0
+            best_error = 1e6
+            for offset in np.linspace(
+                -8 * np.pi, 8 * np.pi, 8 * 4 + 1
+            ):  # intervals of pi/2
+                error = get_error(offset, i, curr_joints)
+                if error < best_error:
+                    best_error = error
+                    best_offset = offset
+            best_offsets.append(best_offset)
+        print()
+        print("best offsets               : ", [f"{x:.3f}" for x in best_offsets])
+        print(
+            "best offsets function of pi: ["
+            + ", ".join([f"{int(np.round(x/(np.pi/2)))}*np.pi/2" for x in best_offsets])
+            + " ]",
+        )
+        if args.gripper:
+            print(
+                "gripper open (degrees)       ",
+                np.rad2deg(driver.get_joints()[-1]) - 0.2,
+            )
+            print(
+                "gripper close (degrees)      ",
+                np.rad2deg(driver.get_joints()[-1]) - 42,
+            )
+
+
+def main(args: Args) -> None:
+    get_config(args)
+
+
+if __name__ == "__main__":
+    main(tyro.cli(Args))

scripts/launch.py

@@ -0,0 +1,39 @@
+import os
+import subprocess
+
+current_file_path = os.path.abspath(__file__)
+
+
+def run_docker_container():
+    user = os.getenv("USER")
+    container_name = f"gello_{user}"
+    gello_path = os.path.abspath(os.path.join(current_file_path, "../../"))
+    volume_mapping = f"{gello_path}:/gello"
+
+    cmd = [
+        "docker",
+        "run",
+        "--runtime=nvidia",
+        "--rm",
+        "--name",
+        container_name,
+        "--privileged",
+        "--volume",
+        volume_mapping,
+        "--volume",
+        "/home/gello:/homefolder",
+        "--net=host",
+        "--volume",
+        "/dev/serial/by-id/:/dev/serial/by-id/",
+        "-it",
+        "gello:latest",
+        "bash",
+        "-c",
+        "pip install -e third_party/DynamixelSDK/python && exec bash",
+    ]
+
+    subprocess.run(cmd)
+
+
+if __name__ == "__main__":
+    run_docker_container()

scripts/visualize_example.py

@@ -0,0 +1,41 @@
+import time
+from pathlib import Path
+
+import mujoco
+import mujoco.viewer
+
+
+def main():
+    _PROJECT_ROOT: Path = Path(__file__).parent.parent
+    _MENAGERIE_ROOT: Path = _PROJECT_ROOT / "third_party" / "mujoco_menagerie"
+    xml = _MENAGERIE_ROOT / "franka_emika_panda" / "panda.xml"
+
+    # xml = _MENAGERIE_ROOT / "universal_robots_ur5e" / "ur5e.xml"
+
+    m = mujoco.MjModel.from_xml_path(xml.as_posix())
+    d = mujoco.MjData(m)
+
+    with mujoco.viewer.launch_passive(m, d) as viewer:
+        # Close the viewer automatically after 30 wall-seconds.
+        while viewer.is_running():
+            step_start = time.time()
+
+            # mj_step can be replaced with code that also evaluates
+            # a policy and applies a control signal before stepping the physics.
+            mujoco.mj_step(m, d)
+
+            # Example modification of a viewer option: toggle contact points every two seconds.
+            with viewer.lock():
+                viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = int(d.time % 2)
+
+            # Pick up changes to the physics state, apply perturbations, update options from GUI.
+            viewer.sync()
+
+            # Rudimentary time keeping, will drift relative to wall clock.
+            time_until_next_step = m.opt.timestep - (time.time() - step_start)
+            if time_until_next_step > 0:
+                time.sleep(time_until_next_step)
+
+
+if __name__ == "__main__":
+    main()