diff --git a/.github/workflows/node-hub-ci-cd.yml b/.github/workflows/node-hub-ci-cd.yml index 910bea0a5..49758275c 100644 --- a/.github/workflows/node-hub-ci-cd.yml +++ b/.github/workflows/node-hub-ci-cd.yml @@ -39,7 +39,16 @@ jobs: steps: - name: Checkout repository if: runner.os == 'Linux' || github.event_name == 'workflow_dispatch' || (github.event_name == 'release' && startsWith(github.ref, 'refs/tags/')) - uses: actions/checkout@v2 + uses: actions/checkout@v4 + with: + submodules: true # Make sure to check out the sub-module + + - name: Update submodule + if: runner.os == 'Linux' + run: | + git submodule update --init --recursive + git submodule update --remote --recursive + - name: Free Disk Space (Ubuntu) uses: jlumbroso/free-disk-space@main if: runner.os == 'Linux' diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 000000000..5b3e5af8b --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "node-hub/dora-rdt-1b/dora_rdt_1b/RoboticsDiffusionTransformer"] + path = node-hub/dora-rdt-1b/dora_rdt_1b/RoboticsDiffusionTransformer + url = https://github.com/thu-ml/RoboticsDiffusionTransformer diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d9709c28a..9fc0486b2 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -12,6 +12,8 @@ You can use the standard `cargo check`, `cargo build`, `cargo run`, and `cargo t To run a command for a specific package only, pass e.g. `--package dora-daemon`. Running a command for the whole workspace is possible by passing `--workspace`. + + ## Continuous Integration (CI) We're using [GitHub Actions](https://github.com/features/actions) to run automated checks on all commits and pull requests. diff --git a/examples/piper/README.md b/examples/piper/README.md new file mode 100644 index 000000000..18a1d125a --- /dev/null +++ b/examples/piper/README.md @@ -0,0 +1,59 @@ +# Getting Started with Tracer + Piper + +# DATA COLLECTION + +## Installation (To do once) + +Make sure to: + +```bash +dora build record.yaml + +# Make sure to install from source pyorbbecksdk + +git clone https://github.com/orbbec/pyorbbecsdk +cd pyorbbecsdk +pip3 install -r requirements.txt +mkdir build +cd build +cmake -Dpybind11_DIR=`pybind11-config --cmakedir` .. +make -j4 +make install +cd .. +pip3 install wheel +python3 setup.py bdist_wheel +pip3 install dist/*.whl + +export PYTHONPATH=$PYTHONPATH:$(pwd)/install/lib/ # Make sure to save this in your .bashrc +cd .. + + +# Install ugv_sdk_py from source +git clone https://github.com/westonrobot/ugv_sdk +cd ugv_sdk +python setup.py build_ext --inplace + +export PYTHONPATH=$PYTHONPATH:$(pwd) # Make sure to save this in your .bashrc +cd .. +``` + +### Your bashrc should contain something like this + +```bash +export PYTHONPATH=$PYTHONPATH:/home/agilex/pyorbbecsdk/install/lib/:/home/agilex/ugv_sdk +``` + +## Setup ( Every boot of the computer ) + +```bash +# Run on Agilex provided computer +source /home/agilex/cobot_magic/Piper_ros_private-ros-noetic/can_config.sh +``` + +## Run + +### For recording episode + +```bash +dora run record.yml +``` diff --git a/examples/piper/arms_camera.yml b/examples/piper/arms_camera.yml new file mode 100644 index 000000000..3677683c3 --- /dev/null +++ b/examples/piper/arms_camera.yml @@ -0,0 +1,74 @@ +nodes: + - id: piper_left + path: /home/agilex/1ms.ai/piper_sdk/dora_piper.py + _unstable_deploy: + machine: piper + inputs: + tick: dora/timer/millis/20 + outputs: + - jointstate + env: + CAN_BUS: can_left + + - id: camera_left + path: /home/agilex/1ms.ai/pyorbbecsdk/examples/color_viewer.py + _unstable_deploy: + machine: piper + inputs: + tick: dora/timer/millis/50 + outputs: + - image + env: + DEVICE_INDEX: 0 + ENCODING: jpeg + + - id: camera_center + path: /home/agilex/1ms.ai/pyorbbecsdk/examples/color_viewer.py + _unstable_deploy: + machine: piper + inputs: + tick: dora/timer/millis/50 + outputs: + - image + env: + DEVICE_INDEX: 1 + ENCODING: jpeg + + - id: camera_right + path: /home/agilex/1ms.ai/pyorbbecsdk/examples/color_viewer.py + _unstable_deploy: + machine: piper + inputs: + tick: dora/timer/millis/50 + outputs: + - image + env: + DEVICE_INDEX: 2 + ENCODING: jpeg + # import opencv as cv + # [cv2.VideoCapture(i) for i in range(12)] + + - id: piper_right + path: /home/agilex/1ms.ai/piper_sdk/dora_piper.py + _unstable_deploy: + machine: piper + inputs: + tick: dora/timer/millis/20 + outputs: + - jointstate + env: + CAN_BUS: can_right + + - id: rerun + path: dora-rerun1 + inputs: + jointstate_piper_left: piper_left/jointstate + jointstate_piper_right: piper_right/jointstate + image_camera_left: camera_left/image + image_camera_center: camera_center/image + image_camera_right: camera_right/image + env: + piper_left_urdf: assets/piper_left.urdf + piper_right_urdf: assets/piper_right.urdf + piper_left_transform: 0 0.2 0 + piper_right_transform: 0 -0.2 0 diff --git a/examples/piper/arms_only.yml b/examples/piper/arms_only.yml new file mode 100644 index 000000000..da5be5249 --- /dev/null +++ b/examples/piper/arms_only.yml @@ -0,0 +1,33 @@ +nodes: + - id: piper_left + path: /home/agilex/1ms.ai/piper_sdk/dora_piper.py + _unstable_deploy: + machine: piper + inputs: + tick: dora/timer/millis/20 + outputs: + - jointstate + env: + CAN_BUS: can_left + + - id: piper_right + path: /home/agilex/1ms.ai/piper_sdk/dora_piper.py + _unstable_deploy: + machine: piper + inputs: + tick: dora/timer/millis/20 + outputs: + - jointstate + env: + CAN_BUS: can_right + + - id: rerun + path: dora-rerun + inputs: + jointstate_piper_left: piper_left/jointstate + jointstate_piper_right: piper_right/jointstate + env: + piper_left_urdf: /home/peter/Documents/work/dora/examples/piper/assets/piper_left.urdf + piper_right_urdf: /home/peter/Documents/work/dora/examples/piper/assets/piper_right.urdf + piper_left_transform: 0 0.3 0 + piper_right_transform: 0 -0.3 0 diff --git a/examples/piper/convert.py b/examples/piper/convert.py new file mode 100644 index 000000000..d2715dcc5 --- /dev/null +++ b/examples/piper/convert.py @@ -0,0 +1,81 @@ +import numpy as np +from scipy.spatial.transform import Rotation as R + + +def convert_quaternion_to_euler(quat): + """ + Convert Quaternion (xyzw) to Euler angles (rpy) + """ + # Normalize + quat = quat / np.linalg.norm(quat) + euler = R.from_quat(quat).as_euler("xyz") + + return euler + + +def convert_euler_to_quaternion(euler): + """ + Convert Euler angles (rpy) to Quaternion (xyzw) + """ + quat = R.from_euler("xyz", euler).as_quat() + + return quat + + +def convert_euler_to_rotation_matrix(euler): + """ + Convert Euler angles (rpy) to rotation matrix (3x3). + """ + quat = R.from_euler("xyz", euler).as_matrix() + + return quat + + +def convert_rotation_matrix_to_euler(rotmat): + """ + Convert rotation matrix (3x3) to Euler angles (rpy). + """ + r = R.from_matrix(rotmat) + euler = r.as_euler("xyz", degrees=False) + + return euler + + +def normalize_vector(v): + v_mag = np.linalg.norm(v, axis=-1, keepdims=True) + v_mag = np.maximum(v_mag, 1e-8) + return v / v_mag + + +def cross_product(u, v): + i = u[:, 1] * v[:, 2] - u[:, 2] * v[:, 1] + j = u[:, 2] * v[:, 0] - u[:, 0] * v[:, 2] + k = u[:, 0] * v[:, 1] - u[:, 1] * v[:, 0] + + out = np.stack((i, j, k), axis=1) + return out + + +def compute_rotation_matrix_from_ortho6d(ortho6d): + x_raw = ortho6d[:, 0:3] + y_raw = ortho6d[:, 3:6] + + x = normalize_vector(x_raw) + z = cross_product(x, y_raw) + z = normalize_vector(z) + y = cross_product(z, x) + + x = x.reshape(-1, 3, 1) + y = y.reshape(-1, 3, 1) + z = z.reshape(-1, 3, 1) + matrix = np.concatenate((x, y, z), axis=2) + return matrix + + +def compute_ortho6d_from_rotation_matrix(matrix): + # The ortho6d represents the first two column vectors a1 and a2 of the + # rotation matrix: [ | , |, | ] + # [ a1, a2, a3] + # [ | , |, | ] + ortho6d = matrix[:, :, :2].transpose(0, 2, 1).reshape(matrix.shape[0], -1) + return ortho6d diff --git a/examples/piper/dummy_inference_2.py b/examples/piper/dummy_inference_2.py new file mode 100644 index 000000000..a90fcfecd --- /dev/null +++ b/examples/piper/dummy_inference_2.py @@ -0,0 +1,126 @@ +from dora import Node + + +import numpy as np +import h5py + +f = h5py.File("data/episode_0.hdf5", "r") + +data = f["action"][:] + + +STATE_VEC_IDX_MAPPING = { + # [0, 10): right arm joint positions + **{"arm_joint_{}_pos".format(i): i for i in range(10)}, + **{"right_arm_joint_{}_pos".format(i): i for i in range(10)}, + # [10, 15): right gripper joint positions + **{"gripper_joint_{}_pos".format(i): i + 10 for i in range(5)}, + **{"right_gripper_joint_{}_pos".format(i): i + 10 for i in range(5)}, + "gripper_open": 10, # alias of right_gripper_joint_0_pos + "right_gripper_open": 10, + # [15, 25): right arm joint velocities + **{"arm_joint_{}_vel".format(i): i + 15 for i in range(10)}, + **{"right_arm_joint_{}_vel".format(i): i + 15 for i in range(10)}, + # [25, 30): right gripper joint velocities + **{"gripper_joint_{}_vel".format(i): i + 25 for i in range(5)}, + **{"right_gripper_joint_{}_vel".format(i): i + 25 for i in range(5)}, + "gripper_open_vel": 25, # alias of right_gripper_joint_0_vel + "right_gripper_open_vel": 25, + # [30, 33): right end effector positions + "eef_pos_x": 30, + "right_eef_pos_x": 30, + "eef_pos_y": 31, + "right_eef_pos_y": 31, + "eef_pos_z": 32, + "right_eef_pos_z": 32, + # [33, 39): right end effector 6D pose + "eef_angle_0": 33, + "right_eef_angle_0": 33, + "eef_angle_1": 34, + "right_eef_angle_1": 34, + "eef_angle_2": 35, + "right_eef_angle_2": 35, + "eef_angle_3": 36, + "right_eef_angle_3": 36, + "eef_angle_4": 37, + "right_eef_angle_4": 37, + "eef_angle_5": 38, + "right_eef_angle_5": 38, + # [39, 42): right end effector velocities + "eef_vel_x": 39, + "right_eef_vel_x": 39, + "eef_vel_y": 40, + "right_eef_vel_y": 40, + "eef_vel_z": 41, + "right_eef_vel_z": 41, + # [42, 45): right end effector angular velocities + "eef_angular_vel_roll": 42, + "right_eef_angular_vel_roll": 42, + "eef_angular_vel_pitch": 43, + "right_eef_angular_vel_pitch": 43, + "eef_angular_vel_yaw": 44, + "right_eef_angular_vel_yaw": 44, + # [45, 50): reserved + # [50, 60): left arm joint positions + **{"left_arm_joint_{}_pos".format(i): i + 50 for i in range(10)}, + # [60, 65): left gripper joint positions + **{"left_gripper_joint_{}_pos".format(i): i + 60 for i in range(5)}, + "left_gripper_open": 60, # alias of left_gripper_joint_0_pos + # [65, 75): left arm joint velocities + **{"left_arm_joint_{}_vel".format(i): i + 65 for i in range(10)}, + # [75, 80): left gripper joint velocities + **{"left_gripper_joint_{}_vel".format(i): i + 75 for i in range(5)}, + "left_gripper_open_vel": 75, # alias of left_gripper_joint_0_vel + # [80, 83): left end effector positions + "left_eef_pos_x": 80, + "left_eef_pos_y": 81, + "left_eef_pos_z": 82, + # [83, 89): left end effector 6D pose + "left_eef_angle_0": 83, + "left_eef_angle_1": 84, + "left_eef_angle_2": 85, + "left_eef_angle_3": 86, + "left_eef_angle_4": 87, + "left_eef_angle_5": 88, + # [89, 92): left end effector velocities + "left_eef_vel_x": 89, + "left_eef_vel_y": 90, + "left_eef_vel_z": 91, + # [92, 95): left end effector angular velocities + "left_eef_angular_vel_roll": 92, + "left_eef_angular_vel_pitch": 93, + "left_eef_angular_vel_yaw": 94, + # [95, 100): reserved + # [100, 102): base linear velocities + "base_vel_x": 100, + "base_vel_y": 101, + # [102, 103): base angular velocities + "base_angular_vel": 102, + # [103, 128): reserved +} + +import time +import pyarrow as pa + +node = Node() +LEFT_UNI_STATE_INDICES = [ + STATE_VEC_IDX_MAPPING[f"left_arm_joint_{i}_pos"] for i in range(6) +] + [STATE_VEC_IDX_MAPPING["left_gripper_open"]] +RIGHT_UNI_STATE_INDICES = [ + STATE_VEC_IDX_MAPPING[f"right_arm_joint_{i}_pos"] for i in range(6) +] + [STATE_VEC_IDX_MAPPING["right_gripper_open"]] +MOBILE_BASE_UNI_STATE_INDICES = [STATE_VEC_IDX_MAPPING["base_vel_x"]] + [ + STATE_VEC_IDX_MAPPING["base_angular_vel"] +] + +for joint in data: + node.send_output( + "jointstate_left", pa.array(joint[LEFT_UNI_STATE_INDICES], type=pa.float32()) + ) + node.send_output( + "jointstate_right", pa.array(joint[RIGHT_UNI_STATE_INDICES], type=pa.float32()) + ) + node.send_output( + "mobile_base", pa.array(joint[MOBILE_BASE_UNI_STATE_INDICES], type=pa.float32()) + ) + time.sleep(0.05) diff --git a/examples/piper/inference.desktop b/examples/piper/inference.desktop new file mode 100644 index 000000000..02086e267 --- /dev/null +++ b/examples/piper/inference.desktop @@ -0,0 +1,7 @@ +[Desktop Entry] +Version=1.0 +Type=Application +Name=Inference +Exec=bash -c "/home/agilex/Desktop/inference.sh; sleep 5" +Terminal=true +Categories=Application diff --git a/examples/piper/inference.sh b/examples/piper/inference.sh new file mode 100644 index 000000000..ee2284f99 --- /dev/null +++ b/examples/piper/inference.sh @@ -0,0 +1,10 @@ +#!/bin/bash +if ! ifconfig | grep -q "can_left"; then + source /home/agilex/cobot_magic/Piper_ros_private-ros-noetic/can_config.sh + sleep 5 +fi +export PYTHONPATH=$PYTHONPATH:/home/agilex/1ms.ai/pyorbbecsdk/install/lib/:/home/agilex/1ms.ai/ugv_sdk +source /home/agilex/miniconda3/etc/profile.d/conda.sh +pkill dora +conda activate dora +dora run /home/agilex/1ms.ai/dora/examples/piper/rdt_1b.yml diff --git a/examples/piper/play_dummy_inference.yml b/examples/piper/play_dummy_inference.yml index d821837c3..0ae5356a1 100644 --- a/examples/piper/play_dummy_inference.yml +++ b/examples/piper/play_dummy_inference.yml @@ -1,13 +1,12 @@ nodes: - id: piper - path: dummy_inference.py + path: dummy_inference_2.py inputs: tick: dora/timer/millis/20 outputs: - jointstate_left - jointstate_right - env: - CAN_BUS: can_left + - mobile_base - id: rerun build: | @@ -20,14 +19,48 @@ nodes: proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/link6.STL proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/link7.STL proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/link8.STL - pip install -e ../../node-hub/dora-rerun pip install git+https://github.com/rerun-io/rerun-loader-python-example-urdf.git + pip install -e ../../node-hub/dora-rerun path: dora-rerun inputs: - jointstate_piper_left: piper/jointstate_left - jointstate_piper_right: piper/jointstate_right + jointstate_piper_left: piper_left/jointstate + jointstate_piper_right: piper_right/jointstate + jointstate_piper_left_pred: piper/jointstate_left + jointstate_piper_right_pred: piper/jointstate_right + series_piper_left: piper_left/jointstate + series_piper_right: piper_right/jointstate + series_piper_left_pred: piper/jointstate_left + series_piper_right_pred: piper/jointstate_right env: piper_left_urdf: piper_left.urdf # Make sure to download meshes from https://github.com/agilexrobotics/Piper_ros/tree/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes and put them in the assets folder piper_right_urdf: piper_right.urdf # Make sure to download meshes from https://github.com/agilexrobotics/Piper_ros/tree/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes and put them in the assets folder piper_left_transform: 0 0.2 0 piper_right_transform: 0 -0.2 0 + piper_left_pred_urdf: assets/piper_left_pred.urdf + piper_right_pred_urdf: assets/piper_right_pred.urdf + piper_left_pred_transform: 0 0.2 0 + piper_right_pred_transform: 0 -0.2 0 + + - id: piper_left + path: /home/agilex/1ms.ai/piper_sdk/dora_piper.py + _unstable_deploy: + machine: piper + inputs: + tick: dora/timer/millis/500 + action: piper/jointstate_left + outputs: + - jointstate + env: + CAN_BUS: can_left + + - id: piper_right + path: /home/agilex/1ms.ai/piper_sdk/dora_piper.py + _unstable_deploy: + machine: piper + inputs: + tick: dora/timer/millis/1000 + action: piper/jointstate_right + outputs: + - jointstate + env: + CAN_BUS: can_right diff --git a/examples/piper/post_process_action.py b/examples/piper/post_process_action.py new file mode 100644 index 000000000..96f27b399 --- /dev/null +++ b/examples/piper/post_process_action.py @@ -0,0 +1,31 @@ +from dora import Node + +node = Node() + + +import time +import pyarrow as pa + +for event in node: + if event["type"] == "INPUT": + actions = event["value"].to_numpy().copy().reshape((64, 14)) + + for action in actions: + gripper_left = action[6] + gripper_right = action[13] + if gripper_right < 0.45: + action[13] = 0.3 + else: + action[13] = 0.6 + + if gripper_left < 0.45: + action[6] = 0.3 + else: + action[6] = 0.6 + + node.send_output("jointstate_left", pa.array(action[:7], type=pa.float32())) + node.send_output( + "jointstate_right", pa.array(action[7:], type=pa.float32()) + ) + time.sleep(0.02) + print(actions) diff --git a/examples/piper/rdt_1b.yml b/examples/piper/rdt_1b.yml new file mode 100644 index 000000000..25c0cabeb --- /dev/null +++ b/examples/piper/rdt_1b.yml @@ -0,0 +1,120 @@ +nodes: + - id: piper_left + build: pip install -e ../../node-hub/dora-piper + path: dora-piper + inputs: + tick: dora/timer/millis/100 + joint_action: post_process_rdt_1b/jointstate_left + outputs: + - jointstate + - pose + - gripper + env: + CAN_BUS: can_left + + - id: piper_right + build: pip install -e ../../node-hub/dora-piper + path: dora-piper + inputs: + tick: dora/timer/millis/100 + joint_action: post_process_rdt_1b/jointstate_right + outputs: + - jointstate + - pose + - gripper + env: + CAN_BUS: can_right + + - id: camera_left + build: pip install -e ../../node-hub/dora-pyorbbecksdk + path: dora-pyorbbecksdk + inputs: + tick: dora/timer/millis/100 + outputs: + - image + env: + DEVICE_INDEX: 0 + ENCODING: jpeg + + - id: camera_center + build: pip install -e ../../node-hub/dora-pyorbbecksdk + path: dora-pyorbbecksdk + inputs: + tick: dora/timer/millis/100 + outputs: + - image + env: + DEVICE_INDEX: 1 + ENCODING: jpeg + + - id: camera_right + build: pip install -e ../../node-hub/dora-pyorbbecksdk + path: dora-pyorbbecksdk + inputs: + tick: dora/timer/millis/100 + outputs: + - image + env: + DEVICE_INDEX: 2 + ENCODING: jpeg + # import opencv as cv + # [cv2.VideoCapture(i) for i in range(12)] + + - id: rerun + path: dora-rerun + inputs: + jointstate_piper_left: piper_left/jointstate + jointstate_piper_right: piper_right/jointstate + jointstate_piper_left_pred: post_process_rdt_1b/jointstate_left + jointstate_piper_right_pred: post_process_rdt_1b/jointstate_right + series_piper_left: piper_left/jointstate + series_piper_right: piper_right/jointstate + series_piper_left_pred: post_process_rdt_1b/jointstate_left + series_piper_right_pred: post_process_rdt_1b/jointstate_right + image_left: camera_left/image + image_center: camera_center/image + image_right: camera_right/image + env: + piper_left_urdf: piper_left.urdf + piper_right_urdf: piper_right.urdf + piper_left_transform: 0 0.2 0 + piper_right_transform: 0 -0.2 0 + piper_left_pred_urdf: piper_left_pred.urdf + piper_right_pred_urdf: piper_right_pred.urdf + piper_left_pred_transform: 0 0.2 0 + piper_right_pred_transform: 0 -0.2 0 + + - id: rdt_1b + path: dora-rdt-1b + inputs: + jointstate_left: + source: piper_left/jointstate + queue_size: 1 + jointstate_right: + source: piper_right/jointstate + queue_size: 1 + image_left: + source: camera_left/image + queue_size: 1 + image_center: + source: camera_center/image + queue_size: 1 + image_right: + source: camera_right/image + queue_size: 1 + tick: + source: dora/timer/secs/2 + queue_size: 1 + outputs: + - action + env: + ROBOTIC_MODEL_NAME_OR_PATH: checkpoints/checkpoint-4400 + LANGUAGE_EMBEDDING_PATH: lang_embed.pt + + - id: post_process_rdt_1b + path: post_process_action.py + inputs: + action: rdt_1b/action + outputs: + - jointstate_left + - jointstate_right diff --git a/examples/piper/record.desktop b/examples/piper/record.desktop new file mode 100644 index 000000000..5c9045817 --- /dev/null +++ b/examples/piper/record.desktop @@ -0,0 +1,7 @@ +[Desktop Entry] +Version=1.0 +Type=Application +Name=Record +Exec=bash -c "/home/agilex/Desktop/record.sh; sleep 5" +Terminal=true +Categories=Application diff --git a/examples/piper/record.py b/examples/piper/record.py new file mode 100644 index 000000000..c12d0b5eb --- /dev/null +++ b/examples/piper/record.py @@ -0,0 +1,274 @@ +import h5py + +import os +import datetime + +from dora import Node +import numpy as np +from convert import ( + convert_euler_to_rotation_matrix, + compute_ortho6d_from_rotation_matrix, +) + +STATE_VEC_IDX_MAPPING = { + # [0, 10): right arm joint positions + **{"arm_joint_{}_pos".format(i): i for i in range(10)}, + **{"right_arm_joint_{}_pos".format(i): i for i in range(10)}, + # [10, 15): right gripper joint positions + **{"gripper_joint_{}_pos".format(i): i + 10 for i in range(5)}, + **{"right_gripper_joint_{}_pos".format(i): i + 10 for i in range(5)}, + "gripper_open": 10, # alias of right_gripper_joint_0_pos + "right_gripper_open": 10, + # [15, 25): right arm joint velocities + **{"arm_joint_{}_vel".format(i): i + 15 for i in range(10)}, + **{"right_arm_joint_{}_vel".format(i): i + 15 for i in range(10)}, + # [25, 30): right gripper joint velocities + **{"gripper_joint_{}_vel".format(i): i + 25 for i in range(5)}, + **{"right_gripper_joint_{}_vel".format(i): i + 25 for i in range(5)}, + "gripper_open_vel": 25, # alias of right_gripper_joint_0_vel + "right_gripper_open_vel": 25, + # [30, 33): right end effector positions + "eef_pos_x": 30, + "right_eef_pos_x": 30, + "eef_pos_y": 31, + "right_eef_pos_y": 31, + "eef_pos_z": 32, + "right_eef_pos_z": 32, + # [33, 39): right end effector 6D pose + "eef_angle_0": 33, + "right_eef_angle_0": 33, + "eef_angle_1": 34, + "right_eef_angle_1": 34, + "eef_angle_2": 35, + "right_eef_angle_2": 35, + "eef_angle_3": 36, + "right_eef_angle_3": 36, + "eef_angle_4": 37, + "right_eef_angle_4": 37, + "eef_angle_5": 38, + "right_eef_angle_5": 38, + # [39, 42): right end effector velocities + "eef_vel_x": 39, + "right_eef_vel_x": 39, + "eef_vel_y": 40, + "right_eef_vel_y": 40, + "eef_vel_z": 41, + "right_eef_vel_z": 41, + # [42, 45): right end effector angular velocities + "eef_angular_vel_roll": 42, + "right_eef_angular_vel_roll": 42, + "eef_angular_vel_pitch": 43, + "right_eef_angular_vel_pitch": 43, + "eef_angular_vel_yaw": 44, + "right_eef_angular_vel_yaw": 44, + # [45, 50): reserved + # [50, 60): left arm joint positions + **{"left_arm_joint_{}_pos".format(i): i + 50 for i in range(10)}, + # [60, 65): left gripper joint positions + **{"left_gripper_joint_{}_pos".format(i): i + 60 for i in range(5)}, + "left_gripper_open": 60, # alias of left_gripper_joint_0_pos + # [65, 75): left arm joint velocities + **{"left_arm_joint_{}_vel".format(i): i + 65 for i in range(10)}, + # [75, 80): left gripper joint velocities + **{"left_gripper_joint_{}_vel".format(i): i + 75 for i in range(5)}, + "left_gripper_open_vel": 75, # alias of left_gripper_joint_0_vel + # [80, 83): left end effector positions + "left_eef_pos_x": 80, + "left_eef_pos_y": 81, + "left_eef_pos_z": 82, + # [83, 89): left end effector 6D pose + "left_eef_angle_0": 83, + "left_eef_angle_1": 84, + "left_eef_angle_2": 85, + "left_eef_angle_3": 86, + "left_eef_angle_4": 87, + "left_eef_angle_5": 88, + # [89, 92): left end effector velocities + "left_eef_vel_x": 89, + "left_eef_vel_y": 90, + "left_eef_vel_z": 91, + # [92, 95): left end effector angular velocities + "left_eef_angular_vel_roll": 92, + "left_eef_angular_vel_pitch": 93, + "left_eef_angular_vel_yaw": 94, + # [95, 100): reserved + # [100, 102): base linear velocities + "base_vel_x": 100, + "base_vel_y": 101, + # [102, 103): base angular velocities + "base_angular_vel": 102, + # [103, 128): reserved +} +STATE_VEC_LEN = 128 + + +now = datetime.datetime.now() + +DATA_DIR = "/home/agilex/Desktop/" + now.strftime("%Y.%m.%d.%H.%M") +os.makedirs(DATA_DIR, exist_ok=True) + +## Make data dir if it does not exist +if not os.path.exists(DATA_DIR): + os.makedirs(DATA_DIR) + + +def save_data(data_dict, dataset_path, data_size): + with h5py.File(dataset_path + ".hdf5", "w", rdcc_nbytes=1024**2 * 2) as root: + root.attrs["sim"] = False + root.attrs["compress"] = False + + obs = root.create_group("observations") + variable_length = h5py.vlen_dtype(np.dtype("uint8")) + image = obs.create_group("images") + _ = image.create_dataset( + "cam_high", + (data_size,), + dtype=variable_length, + ) + _ = image.create_dataset( + "cam_left_wrist", + (data_size,), + dtype=variable_length, + ) + _ = image.create_dataset( + "cam_right_wrist", + (data_size,), + dtype=variable_length, + ) + + _ = obs.create_dataset("qpos", (data_size, 128)) + _ = root.create_dataset("action", (data_size, 128)) + + # data_dict write into h5py.File + for name, array in data_dict.items(): + print(name) + if "images" in name: + image[name][...] = array + else: + root[name][...] = array + + +data_dict = { + "/observations/qpos": [], + "/observations/images/cam_high": [], + "/observations/images/cam_left_wrist": [], + "/observations/images/cam_right_wrist": [], + "/action": [], +} + + +node = Node() + +LEAD_CAMERA = "/observations/images/cam_high" + +tmp_dict = {} + +i = 0 + +start = False +for event in node: + if event["type"] == "INPUT": + if "save" in event["id"]: + char = event["value"][0].as_py() + if char == "p": + if start == False: + continue + + save_data( + data_dict, + f"{DATA_DIR}/episode_{i}", + len(data_dict["/observations/qpos"]), + ) + + # Reset dict + data_dict = { + "/observations/qpos": [], + "/observations/images/cam_high": [], + "/observations/images/cam_left_wrist": [], + "/observations/images/cam_right_wrist": [], + "/action": [], + } + i += 1 + start = False + elif char == "s": + start = True + + elif "image" in event["id"]: + tmp_dict[event["id"]] = event["value"].to_numpy() + elif "qpos" in event["id"]: + tmp_dict[event["id"]] = event["value"].to_numpy() + elif "pose" in event["id"]: + value = event["value"].to_numpy() + euler = value[None, 3:6] # Add batch dimension + rotmat = convert_euler_to_rotation_matrix(euler) + ortho6d = compute_ortho6d_from_rotation_matrix(rotmat)[0] + values = np.array( + [ + value[0], + value[1], + value[2], + ortho6d[0], + ortho6d[1], + ortho6d[2], + ortho6d[3], + ortho6d[4], + ortho6d[5], + ] + ) + tmp_dict[event["id"]] = values + elif "base_vel" in event["id"]: + tmp_dict[event["id"]] = event["value"].to_numpy() + + # Check if tmp dict is full + if len(tmp_dict) != 7: + continue + elif event["id"] == LEAD_CAMERA and start == True: + values = np.concatenate( + [ + tmp_dict["/observations/qpos_left"], + tmp_dict["/observations/qpos_right"], + tmp_dict["/observations/pose_left"], + tmp_dict["/observations/pose_right"], + # tmp_dict["/observations/base_vel"], + ] + ) + UNI_STATE_INDICES = ( + [STATE_VEC_IDX_MAPPING[f"left_arm_joint_{i}_pos"] for i in range(6)] + + [STATE_VEC_IDX_MAPPING["left_gripper_open"]] + + [STATE_VEC_IDX_MAPPING[f"right_arm_joint_{i}_pos"] for i in range(6)] + + [STATE_VEC_IDX_MAPPING["right_gripper_open"]] + + [STATE_VEC_IDX_MAPPING["left_eef_pos_x"]] + + [STATE_VEC_IDX_MAPPING["left_eef_pos_y"]] + + [STATE_VEC_IDX_MAPPING["left_eef_pos_z"]] + + [STATE_VEC_IDX_MAPPING["left_eef_angle_0"]] + + [STATE_VEC_IDX_MAPPING["left_eef_angle_1"]] + + [STATE_VEC_IDX_MAPPING["left_eef_angle_2"]] + + [STATE_VEC_IDX_MAPPING["left_eef_angle_3"]] + + [STATE_VEC_IDX_MAPPING["left_eef_angle_4"]] + + [STATE_VEC_IDX_MAPPING["left_eef_angle_5"]] + + [STATE_VEC_IDX_MAPPING["right_eef_pos_x"]] + + [STATE_VEC_IDX_MAPPING["right_eef_pos_y"]] + + [STATE_VEC_IDX_MAPPING["right_eef_pos_z"]] + + [STATE_VEC_IDX_MAPPING["right_eef_angle_0"]] + + [STATE_VEC_IDX_MAPPING["right_eef_angle_1"]] + + [STATE_VEC_IDX_MAPPING["right_eef_angle_2"]] + + [STATE_VEC_IDX_MAPPING["right_eef_angle_3"]] + + [STATE_VEC_IDX_MAPPING["right_eef_angle_4"]] + + [STATE_VEC_IDX_MAPPING["right_eef_angle_5"]] + # + [STATE_VEC_IDX_MAPPING["base_vel_x"]] + # + [STATE_VEC_IDX_MAPPING["base_angular_vel"]], + ) + universal_vec = np.zeros(STATE_VEC_LEN) + universal_vec[UNI_STATE_INDICES] = values + data_dict["/observations/qpos"].append(universal_vec) + # We reproduce obs and action + data_dict["/action"].append(universal_vec) + data_dict["/observations/images/cam_high"].append( + tmp_dict["/observations/images/cam_high"] + ) + data_dict["/observations/images/cam_left_wrist"].append( + tmp_dict["/observations/images/cam_left_wrist"] + ) + data_dict["/observations/images/cam_right_wrist"].append( + tmp_dict["/observations/images/cam_right_wrist"] + ) diff --git a/examples/piper/record.sh b/examples/piper/record.sh new file mode 100644 index 000000000..190dccd6e --- /dev/null +++ b/examples/piper/record.sh @@ -0,0 +1,10 @@ +#!/bin/bash +if ! ifconfig | grep -q "can_left"; then + source /home/agilex/cobot_magic/Piper_ros_private-ros-noetic/can_config.sh + sleep 5 +fi +export PYTHONPATH=$PYTHONPATH:/home/agilex/1ms.ai/pyorbbecsdk/install/lib/:/home/agilex/1ms.ai/ugv_sdk +source /home/agilex/miniconda3/etc/profile.d/conda.sh +pkill dora +conda activate dora +dora run /home/agilex/1ms.ai/dora/examples/piper/record.yml diff --git a/examples/piper/record.yml b/examples/piper/record.yml new file mode 100644 index 000000000..d890c5a1d --- /dev/null +++ b/examples/piper/record.yml @@ -0,0 +1,114 @@ +nodes: + - id: piper_left + build: pip install -e ../../node-hub/dora-piper + path: dora-piper + inputs: + tick: dora/timer/millis/40 + outputs: + - jointstate + - pose + - gripper + env: + CAN_BUS: can_left + TEACH_MODE: True + + - id: piper_right + build: pip install -e ../../node-hub/dora-piper + path: dora-piper + inputs: + tick: dora/timer/millis/40 + outputs: + - jointstate + - pose + - gripper + env: + CAN_BUS: can_right + TEACH_MODE: True + + - id: camera_left + build: pip install -e ../../node-hub/dora-pyorbbecksdk + path: dora-pyorbbecksdk + inputs: + tick: dora/timer/millis/40 + outputs: + - image + env: + DEVICE_INDEX: 0 + ENCODING: jpeg + + - id: camera_center + build: pip install -e ../../node-hub/dora-pyorbbecksdk + path: dora-pyorbbecksdk + inputs: + tick: dora/timer/millis/40 + outputs: + - image + env: + DEVICE_INDEX: 1 + ENCODING: jpeg + + - id: camera_right + build: pip install -e ../../node-hub/dora-pyorbbecksdk + path: dora-pyorbbecksdk + inputs: + tick: dora/timer/millis/40 + outputs: + - image + env: + DEVICE_INDEX: 2 + ENCODING: jpeg + + - id: rerun + path: dora-rerun + build: | + proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/base_link.STL + proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/link1.STL + proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/link2.STL + proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/link3.STL + proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/link4.STL + proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/link5.STL + proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/link6.STL + proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/link7.STL + proxychains wget -nc https://raw.githubusercontent.com/agilexrobotics/Piper_ros/4f22c61f96b8fb3ef3f937b99b63edb697caadf0/src/piper_description/meshes/link8.STL + pip install git+https://github.com/rerun-io/rerun-loader-python-example-urdf.git + pip install dora-rerun + inputs: + jointstate_piper_left: piper_left/jointstate + jointstate_piper_right: piper_right/jointstate + image_left: camera_left/image + image_center: camera_center/image + image_right: camera_right/image + env: + piper_left_urdf: piper_left.urdf + piper_right_urdf: piper_right.urdf + piper_left_transform: 0 0.2 0 + piper_right_transform: 0 -0.2 0 + + - id: keyboard + build: pip install dora-keyboard + path: dora-keyboard + inputs: + tick: dora/timer/millis/1000 + outputs: + - char + env: + DISPLAY: :0 + + - id: recorder + path: record.py + inputs: + /observations/qpos_left: + source: piper_left/jointstate + /observations/qpos_right: + source: piper_right/jointstate + /observations/pose_left: + source: piper_left/pose + /observations/pose_right: + source: piper_right/pose + /observations/images/cam_left_wrist: + source: camera_left/image + /observations/images/cam_high: + source: camera_center/image + /observations/images/cam_right_wrist: + source: camera_right/image + save: keyboard/char diff --git a/examples/piper/replay.py b/examples/piper/replay.py new file mode 100644 index 000000000..24bc1618e --- /dev/null +++ b/examples/piper/replay.py @@ -0,0 +1,129 @@ +from dora import Node + + +import numpy as np +import h5py +import os + +EPISODE_PATH = os.getenv("EPISODE_PATH", "data/episode_0.hdf5") + +f = h5py.File(EPISODE_PATH, "r") + +data = f["action"][:] + + +STATE_VEC_IDX_MAPPING = { + # [0, 10): right arm joint positions + **{"arm_joint_{}_pos".format(i): i for i in range(10)}, + **{"right_arm_joint_{}_pos".format(i): i for i in range(10)}, + # [10, 15): right gripper joint positions + **{"gripper_joint_{}_pos".format(i): i + 10 for i in range(5)}, + **{"right_gripper_joint_{}_pos".format(i): i + 10 for i in range(5)}, + "gripper_open": 10, # alias of right_gripper_joint_0_pos + "right_gripper_open": 10, + # [15, 25): right arm joint velocities + **{"arm_joint_{}_vel".format(i): i + 15 for i in range(10)}, + **{"right_arm_joint_{}_vel".format(i): i + 15 for i in range(10)}, + # [25, 30): right gripper joint velocities + **{"gripper_joint_{}_vel".format(i): i + 25 for i in range(5)}, + **{"right_gripper_joint_{}_vel".format(i): i + 25 for i in range(5)}, + "gripper_open_vel": 25, # alias of right_gripper_joint_0_vel + "right_gripper_open_vel": 25, + # [30, 33): right end effector positions + "eef_pos_x": 30, + "right_eef_pos_x": 30, + "eef_pos_y": 31, + "right_eef_pos_y": 31, + "eef_pos_z": 32, + "right_eef_pos_z": 32, + # [33, 39): right end effector 6D pose + "eef_angle_0": 33, + "right_eef_angle_0": 33, + "eef_angle_1": 34, + "right_eef_angle_1": 34, + "eef_angle_2": 35, + "right_eef_angle_2": 35, + "eef_angle_3": 36, + "right_eef_angle_3": 36, + "eef_angle_4": 37, + "right_eef_angle_4": 37, + "eef_angle_5": 38, + "right_eef_angle_5": 38, + # [39, 42): right end effector velocities + "eef_vel_x": 39, + "right_eef_vel_x": 39, + "eef_vel_y": 40, + "right_eef_vel_y": 40, + "eef_vel_z": 41, + "right_eef_vel_z": 41, + # [42, 45): right end effector angular velocities + "eef_angular_vel_roll": 42, + "right_eef_angular_vel_roll": 42, + "eef_angular_vel_pitch": 43, + "right_eef_angular_vel_pitch": 43, + "eef_angular_vel_yaw": 44, + "right_eef_angular_vel_yaw": 44, + # [45, 50): reserved + # [50, 60): left arm joint positions + **{"left_arm_joint_{}_pos".format(i): i + 50 for i in range(10)}, + # [60, 65): left gripper joint positions + **{"left_gripper_joint_{}_pos".format(i): i + 60 for i in range(5)}, + "left_gripper_open": 60, # alias of left_gripper_joint_0_pos + # [65, 75): left arm joint velocities + **{"left_arm_joint_{}_vel".format(i): i + 65 for i in range(10)}, + # [75, 80): left gripper joint velocities + **{"left_gripper_joint_{}_vel".format(i): i + 75 for i in range(5)}, + "left_gripper_open_vel": 75, # alias of left_gripper_joint_0_vel + # [80, 83): left end effector positions + "left_eef_pos_x": 80, + "left_eef_pos_y": 81, + "left_eef_pos_z": 82, + # [83, 89): left end effector 6D pose + "left_eef_angle_0": 83, + "left_eef_angle_1": 84, + "left_eef_angle_2": 85, + "left_eef_angle_3": 86, + "left_eef_angle_4": 87, + "left_eef_angle_5": 88, + # [89, 92): left end effector velocities + "left_eef_vel_x": 89, + "left_eef_vel_y": 90, + "left_eef_vel_z": 91, + # [92, 95): left end effector angular velocities + "left_eef_angular_vel_roll": 92, + "left_eef_angular_vel_pitch": 93, + "left_eef_angular_vel_yaw": 94, + # [95, 100): reserved + # [100, 102): base linear velocities + "base_vel_x": 100, + "base_vel_y": 101, + # [102, 103): base angular velocities + "base_angular_vel": 102, + # [103, 128): reserved +} + +import time +import pyarrow as pa + +node = Node() +LEFT_UNI_STATE_INDICES = [ + STATE_VEC_IDX_MAPPING[f"left_arm_joint_{i}_pos"] for i in range(6) +] + [STATE_VEC_IDX_MAPPING["left_gripper_open"]] +RIGHT_UNI_STATE_INDICES = [ + STATE_VEC_IDX_MAPPING[f"right_arm_joint_{i}_pos"] for i in range(6) +] + [STATE_VEC_IDX_MAPPING["right_gripper_open"]] +MOBILE_BASE_UNI_STATE_INDICES = [STATE_VEC_IDX_MAPPING["base_vel_x"]] + [ + STATE_VEC_IDX_MAPPING["base_angular_vel"] +] + +for joint in data: + node.send_output( + "jointstate_left", pa.array(joint[LEFT_UNI_STATE_INDICES], type=pa.float32()) + ) + node.send_output( + "jointstate_right", pa.array(joint[RIGHT_UNI_STATE_INDICES], type=pa.float32()) + ) + # node.send_output( + # "mobile_base", pa.array(joint[MOBILE_BASE_UNI_STATE_INDICES], type=pa.float32()) + # ) + time.sleep(0.05) diff --git a/examples/piper/replay.yml b/examples/piper/replay.yml new file mode 100644 index 000000000..d8df7ead4 --- /dev/null +++ b/examples/piper/replay.yml @@ -0,0 +1,55 @@ +nodes: + - id: piper + path: replay.py + inputs: + tick: dora/timer/millis/20 + outputs: + - jointstate_left + - jointstate_right + - mobile_base + + - id: rerun + build: wget + path: dora-rerun + inputs: + jointstate_piper_left: piper_left/jointstate + jointstate_piper_right: piper_right/jointstate + jointstate_piper_left_pred: piper/jointstate_left + jointstate_piper_right_pred: piper/jointstate_right + series_piper_left: piper_left/jointstate + series_piper_right: piper_right/jointstate + series_piper_left_pred: piper/jointstate_left + series_piper_right_pred: piper/jointstate_right + env: + piper_left_urdf: piper_left.urdf + piper_right_urdf: piper_right.urdf + piper_left_transform: 0 0.2 0 + piper_right_transform: 0 -0.2 0 + piper_left_pred_urdf: piper_left_pred.urdf + piper_right_pred_urdf: piper_right_pred.urdf + piper_left_pred_transform: 0 0.2 0 + piper_right_pred_transform: 0 -0.2 0 + + - id: piper_left + path: /home/agilex/1ms.ai/piper_sdk/dora_piper.py + _unstable_deploy: + machine: piper + inputs: + tick: dora/timer/millis/500 + action: piper/jointstate_left + outputs: + - jointstate + env: + CAN_BUS: can_left + + - id: piper_right + path: /home/agilex/1ms.ai/piper_sdk/dora_piper.py + _unstable_deploy: + machine: piper + inputs: + tick: dora/timer/millis/1000 + action: piper/jointstate_right + outputs: + - jointstate + env: + CAN_BUS: can_right diff --git a/node-hub/dora-keyboard/dora_keyboard/main.py b/node-hub/dora-keyboard/dora_keyboard/main.py index 6e8662637..620f89371 100644 --- a/node-hub/dora-keyboard/dora_keyboard/main.py +++ b/node-hub/dora-keyboard/dora_keyboard/main.py @@ -6,8 +6,14 @@ def main(): node = Node() + + always_none = node.next(timeout=0.001) is None with keyboard.Events() as events: while True: + if not always_none: + event_stream_is_none = node.next(timeout=0.001) is None + if event_stream_is_none: + break event = events.get(1.0) if event is not None and isinstance(event, Events.Press): if hasattr(event.key, "char"): diff --git a/node-hub/dora-piper/dora_piper/main.py b/node-hub/dora-piper/dora_piper/main.py index 0a7e99676..dce1ac4ef 100644 --- a/node-hub/dora-piper/dora_piper/main.py +++ b/node-hub/dora-piper/dora_piper/main.py @@ -1,6 +1,7 @@ from piper_sdk import C_PiperInterface from dora import Node import pyarrow as pa +import numpy as np import os import time @@ -63,22 +64,9 @@ def main(): for event in node: if event["type"] == "INPUT": - if event["id"] != "action": - joint = piper.GetArmJointMsgs() - gripper = piper.GetArmGripperMsgs() - - joint_value = [] - joint_value += [joint.joint_state.joint_1.real / factor] - joint_value += [joint.joint_state.joint_2.real / factor] - joint_value += [joint.joint_state.joint_3.real / factor] - joint_value += [joint.joint_state.joint_4.real / factor] - joint_value += [joint.joint_state.joint_5.real / factor] - joint_value += [joint.joint_state.joint_6.real / factor] - joint_value += [gripper.gripper_state.grippers_angle / 1000 / 100] - - node.send_output("jointstate", pa.array(joint_value, type=pa.float32())) - else: - + if event["id"] == "joint_action": + if TEACH_MODE: + continue # Do not push to many commands to fast. Limiting it to 20Hz if time.time() - elapsed_time > 0.05: elapsed_time = time.time() @@ -98,6 +86,63 @@ def main(): piper.JointCtrl(joint_0, joint_1, joint_2, joint_3, joint_4, joint_5) piper.GripperCtrl(abs(joint_6), 1000, 0x01, 0) piper.MotionCtrl_2(0x01, 0x01, 50, 0x00) + + elif event["id"] == "eef_action": + if TEACH_MODE: + continue + # Do not push to many commands to fast. Limiting it to 20Hz + if time.time() - elapsed_time > 0.05: + elapsed_time = time.time() + else: + continue + + position = event["value"].to_numpy() + piper.MotionCtrl_2(0x01, 0x01, 50, 0x00) + piper.EndPoseCtrl( + position[0] * 1000 * 1000, + position[1] * 1000 * 1000, + position[2] * 1000 * 1000, + position[3] * 1000 / (2 * np.pi) * 360, + position[4] * 1000 / (2 * np.pi) * 360, + position[5] * 1000 / (2 * np.pi) * 360, + ) + piper.GripperCtrl(abs(position[6] * 1000 * 100), 1000, 0x01, 0) + piper.MotionCtrl_2(0x01, 0x01, 50, 0x00) + + else: + joint = piper.GetArmJointMsgs() + + joint_value = [] + joint_value += [joint.joint_state.joint_1.real / factor] + joint_value += [joint.joint_state.joint_2.real / factor] + joint_value += [joint.joint_state.joint_3.real / factor] + joint_value += [joint.joint_state.joint_4.real / factor] + joint_value += [joint.joint_state.joint_5.real / factor] + joint_value += [joint.joint_state.joint_6.real / factor] + + gripper = piper.GetArmGripperMsgs() + joint_value += [gripper.gripper_state.grippers_angle / 1000 / 100] + + node.send_output("jointstate", pa.array(joint_value, type=pa.float32())) + + position = piper.GetArmEndPoseMsgs() + position_value = [] + position_value += [position.end_pose.X_axis * 0.001 * 0.001] + position_value += [position.end_pose.Y_axis * 0.001 * 0.001] + position_value += [position.end_pose.Z_axis * 0.001 * 0.001] + position_value += [position.end_pose.RX_axis * 0.001 / 360 * 2 * np.pi] + position_value += [position.end_pose.RY_axis * 0.001 / 360 * 2 * np.pi] + position_value += [position.end_pose.RZ_axis * 0.001 / 360 * 2 * np.pi] + + node.send_output("pose", pa.array(position_value, type=pa.float32())) + node.send_output( + "gripper", + pa.array( + [gripper.gripper_state.grippers_angle / 1000 / 100], + type=pa.float32(), + ), + ) + elif event["type"] == "STOP": if not TEACH_MODE: diff --git a/node-hub/dora-piper/pyproject.toml b/node-hub/dora-piper/pyproject.toml index 530874193..701abd951 100644 --- a/node-hub/dora-piper/pyproject.toml +++ b/node-hub/dora-piper/pyproject.toml @@ -11,7 +11,10 @@ packages = [{ include = "dora_piper" }] dora-rs = "^0.3.6" python = "^3.7" piper_sdk = "^0.0.8" +numpy = "< 2.0.0" +[tool.poetry.dev-dependencies] +pytest = "^8.3.4" [tool.poetry.scripts] dora-piper = "dora_piper.main:main" diff --git a/node-hub/dora-pyrealsense/pyproject.toml b/node-hub/dora-pyrealsense/pyproject.toml index 383f846c0..88f6f56ef 100644 --- a/node-hub/dora-pyrealsense/pyproject.toml +++ b/node-hub/dora-pyrealsense/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "dora-pyrealsense" -version = "0.0.1" +version = "0.3.8" authors = ["Haixuan Xavier Tao "] description = "Dora Node for capturing video with Pyrealsense" readme = "README.md" diff --git a/node-hub/dora-rdt-1b/README.md b/node-hub/dora-rdt-1b/README.md new file mode 100644 index 000000000..6f2da50fe --- /dev/null +++ b/node-hub/dora-rdt-1b/README.md @@ -0,0 +1,3 @@ +# Dora RDT-1B node + +Experimental node for using a RDT-1B VLA model. diff --git a/node-hub/dora-rdt-1b/dora_rdt_1b/RoboticsDiffusionTransformer b/node-hub/dora-rdt-1b/dora_rdt_1b/RoboticsDiffusionTransformer new file mode 160000 index 000000000..b2889e65c --- /dev/null +++ b/node-hub/dora-rdt-1b/dora_rdt_1b/RoboticsDiffusionTransformer @@ -0,0 +1 @@ +Subproject commit b2889e65cfe62571ced3ce88f00e7d80b41fee69 diff --git a/node-hub/dora-rdt-1b/dora_rdt_1b/__init__.py b/node-hub/dora-rdt-1b/dora_rdt_1b/__init__.py new file mode 100644 index 000000000..ed4e21919 --- /dev/null +++ b/node-hub/dora-rdt-1b/dora_rdt_1b/__init__.py @@ -0,0 +1,19 @@ +import os +import sys +from pathlib import Path + +# Define the path to the README file relative to the package directory +readme_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "README.md") + +# Read the content of the README file +try: + with open(readme_path, "r", encoding="utf-8") as f: + __doc__ = f.read() +except FileNotFoundError: + __doc__ = "README file not found." + + +# Set up the import hook + +submodule_path = Path(__file__).resolve().parent / "RoboticsDiffusionTransformer" +sys.path.insert(0, str(submodule_path)) diff --git a/node-hub/dora-rdt-1b/dora_rdt_1b/main.py b/node-hub/dora-rdt-1b/dora_rdt_1b/main.py new file mode 100644 index 000000000..35f0b7ca0 --- /dev/null +++ b/node-hub/dora-rdt-1b/dora_rdt_1b/main.py @@ -0,0 +1,303 @@ +# install dependencies as shown in the README here https://github.com/alik-git/RoboticsDiffusionTransformer?tab=readme-ov-file#installation +import yaml +import torch +import numpy as np +from PIL import Image + +from dora_rdt_1b.RoboticsDiffusionTransformer.configs.state_vec import ( + STATE_VEC_IDX_MAPPING, +) +from dora import Node +import cv2 +import pyarrow as pa +import os +from pathlib import Path + +VISION_DEFAULT_PATH = "robotics-diffusion-transformer/rdt-1b" +ROBOTIC_MODEL_NAME_OR_PATH = os.getenv( + "ROBOTIC_MODEL_NAME_OR_PATH", VISION_DEFAULT_PATH +) +LANGUAGE_EMBEDDING_PATH = os.getenv("LANGUAGE_EMBEDDING", "lang_embed.pt") + +VISION_DEFAULT_PATH = "google/siglip-so400m-patch14-384" +VISION_MODEL_NAME_OR_PATH = os.getenv("VISION_MODEL_NAME_OR_PATH", VISION_DEFAULT_PATH) + +DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu" +DEVICE = os.getenv("DEVICE", DEVICE) +DEVICE = torch.device(DEVICE) +DTYPE = torch.float16 if DEVICE != "cpu" else torch.float32 + +file_path = Path(__file__).parent + +config_path = ( + file_path / "RoboticsDiffusionTransformer/configs/base.yaml" +) # default config + +with open(config_path, "r", encoding="utf-8") as fp: + config = yaml.safe_load(fp) + + +def get_policy(): + from dora_rdt_1b.RoboticsDiffusionTransformer.models.rdt_runner import RDTRunner + + pretrained_model_name_or_path = ROBOTIC_MODEL_NAME_OR_PATH + rdt = RDTRunner.from_pretrained(pretrained_model_name_or_path) + device = torch.device("cuda:0") + dtype = torch.bfloat16 # recommended + rdt.to(device, dtype=dtype) + rdt.eval() + return rdt + + +def get_vision_model(): + from dora_rdt_1b.RoboticsDiffusionTransformer.models.multimodal_encoder.siglip_encoder import ( + SiglipVisionTower, + ) + + # Load vision encoder + vision_encoder = SiglipVisionTower( + vision_tower=VISION_MODEL_NAME_OR_PATH, + args=None, + ) + vision_encoder.to(DEVICE, dtype=DTYPE) + vision_encoder.eval() + image_processor = vision_encoder.image_processor + return vision_encoder, image_processor + + +def get_language_embeddings(): + device = torch.device("cuda:0") + + lang_embeddings = torch.load( + LANGUAGE_EMBEDDING_PATH, + map_location=device, + ) + + return lang_embeddings.unsqueeze( + 0 + ) # Size: (B, L_lang, D) or None, language condition tokens (variable length), dimension D is assumed to be the same as the hidden size. + + +def expand2square(pil_img, background_color): + width, height = pil_img.size + if width == height: + return pil_img + elif width > height: + result = Image.new(pil_img.mode, (width, width), background_color) + result.paste(pil_img, (0, (width - height) // 2)) + return result + else: + result = Image.new(pil_img.mode, (height, height), background_color) + result.paste(pil_img, ((height - width) // 2, 0)) + return result + + +def process_image(rgbs_lst, image_processor, vision_encoder): + # previous_image_path = "/mnt/hpfs/1ms.ai/dora/node-hub/dora-rdt-1b/dora_rdt_1b/RoboticsDiffusionTransformer/img.jpeg" + # # previous_image = None # if t = 0 + # previous_image = Image.fromarray(previous_image_path).convert("RGB") # if t > 0 + + # current_image_path = "/mnt/hpfs/1ms.ai/dora/node-hub/dora-rdt-1b/dora_rdt_1b/RoboticsDiffusionTransformer/img.jpeg" + # current_image = Image.fromarray(current_image_path).convert("RGB") + + # here I suppose you only have an image from exterior (e.g., 3rd person view) and you don't have any state information + # the images should arrange in sequence [exterior_image, right_wrist_image, left_wrist_image] * image_history_size (e.g., 2) + # rgbs_lst = [[previous_image, None, None], [current_image, None, None]] + # if your have an right_wrist_image, then it should be + # rgbs_lst = [ + # [previous_image, previous_right_wrist_image, None], + # [current_image, current_right_wrist_image, None] + # ] + + # image pre-processing + # The background image used for padding + + image_tensor_list = [] + for step in range(config["common"]["img_history_size"]): + rgbs = rgbs_lst[step] + for rgb in rgbs: + assert rgb, "You should not have None image" + image = rgb + + if config["dataset"].get("image_aspect_ratio", "pad") == "pad": + background_color = tuple( + int(x * 255) for x in image_processor.image_mean + ) + image = expand2square(image, background_color) + image = image_processor.preprocess(image, return_tensors="pt")[ + "pixel_values" + ][0] + image_tensor_list.append(image) + + image_tensor = torch.stack(image_tensor_list, dim=0).to(DEVICE, dtype=DTYPE) + # encode images + image_embeds = vision_encoder(image_tensor).detach() + return image_embeds.reshape(-1, vision_encoder.hidden_size).unsqueeze(0) + + +def get_states(proprio): + # suppose you control in 7DOF joint position + STATE_INDICES = [ + STATE_VEC_IDX_MAPPING["left_arm_joint_0_pos"], + STATE_VEC_IDX_MAPPING["left_arm_joint_1_pos"], + STATE_VEC_IDX_MAPPING["left_arm_joint_2_pos"], + STATE_VEC_IDX_MAPPING["left_arm_joint_3_pos"], + STATE_VEC_IDX_MAPPING["left_arm_joint_4_pos"], + STATE_VEC_IDX_MAPPING["left_arm_joint_5_pos"], + STATE_VEC_IDX_MAPPING["left_gripper_open"], + STATE_VEC_IDX_MAPPING["right_arm_joint_0_pos"], + STATE_VEC_IDX_MAPPING["right_arm_joint_1_pos"], + STATE_VEC_IDX_MAPPING["right_arm_joint_2_pos"], + STATE_VEC_IDX_MAPPING["right_arm_joint_3_pos"], + STATE_VEC_IDX_MAPPING["right_arm_joint_4_pos"], + STATE_VEC_IDX_MAPPING["right_arm_joint_5_pos"], + STATE_VEC_IDX_MAPPING["right_gripper_open"], + ] + + B, N = 1, 1 # batch size and state history size + states = torch.zeros( + (B, N, config["model"]["state_token_dim"]), device=DEVICE, dtype=DTYPE + ) + # suppose you do not have proprio + # it's kind of tricky, I strongly suggest adding proprio as input and further fine-tuning + proprio = torch.tensor(proprio, device=DEVICE, dtype=DTYPE).reshape( + (1, 1, -1) + ) # B, N = 1, 1 # batch size and state history size + + # if you have proprio, you can do like this + # format like this: [arm_joint_0_pos, arm_joint_1_pos, arm_joint_2_pos, arm_joint_3_pos, arm_joint_4_pos, arm_joint_5_pos, arm_joint_6_pos, gripper_open] + # proprio = torch.tensor([0, 1, 2, 3, 4, 5, 6, 0.5]).reshape((1, 1, -1)) + states[:, :, STATE_INDICES] = proprio + + state_elem_mask = torch.zeros( + (1, config["model"]["state_token_dim"]), device=DEVICE, dtype=torch.bool + ) + + state_elem_mask[:, STATE_INDICES] = True + states, state_elem_mask = states.to(DEVICE, dtype=DTYPE), state_elem_mask.to( + DEVICE, dtype=DTYPE + ) + states = states[:, -1:, :] # only use the last state + return states, state_elem_mask, STATE_INDICES + + +def main(): + + rdt = get_policy() + lang_embeddings = get_language_embeddings() + vision_encoder, image_processor = get_vision_model() + + ## for image + # image_embeds = process_image(rgb_lst, image_processor, vision_encoder) + ## for states + # states, state_elem_mask, STATE_INDICES = get_states(states) + node = Node() + frames = {} + joints = {} + with torch.no_grad(): + + for event in node: + event_type = event["type"] + if event_type == "INPUT": + + event_id = event["id"] + + if "image" in event_id: + storage = event["value"] + metadata = event["metadata"] + encoding = metadata["encoding"] + + if encoding == "bgr8": + channels = 3 + storage_type = np.uint8 + elif encoding == "rgb8": + channels = 3 + storage_type = np.uint8 + elif encoding in ["jpeg", "jpg", "jpe", "bmp", "webp", "png"]: + channels = 3 + storage_type = np.uint8 + else: + raise RuntimeError(f"Unsupported image encoding: {encoding}") + + if encoding == "bgr8": + width = metadata["width"] + height = metadata["height"] + frame = ( + storage.to_numpy() + .astype(storage_type) + .reshape((height, width, channels)) + ) + frame = frame[:, :, ::-1] # OpenCV image (BGR to RGB) + elif encoding == "rgb8": + width = metadata["width"] + height = metadata["height"] + frame = ( + storage.to_numpy() + .astype(storage_type) + .reshape((height, width, channels)) + ) + elif encoding in ["jpeg", "jpg", "jpe", "bmp", "webp", "png"]: + storage = storage.to_numpy() + frame = cv2.imdecode(storage, cv2.IMREAD_COLOR) + frame = frame[:, :, ::-1] # OpenCV image (BGR to RGB) + else: + raise RuntimeError(f"Unsupported image encoding: {encoding}") + frames[f"last_{event_id}"] = frames.get( + event_id, Image.fromarray(frame) + ) + frames[event_id] = Image.fromarray(frame) + elif "jointstate" in event_id: + joints[event_id] = event["value"].to_numpy() + + elif "tick" == event_id: + ## Wait for all images + if len(frames.keys()) < 6: + continue + if len(joints.keys()) < 2: + continue + + ## Embed images + rgbs_lst = [ + [ + frames["last_image_center"], + frames["last_image_right"], + frames["last_image_left"], + ], + [ + frames["image_center"], + frames["image_right"], + frames["image_left"], + ], + ] + image_embeds = process_image( + rgbs_lst, image_processor, vision_encoder + ) + + ## Embed states + proprio = np.concatenate( + [ + joints["jointstate_left"], + joints["jointstate_right"], + ] + ) + states, state_elem_mask, state_indices = get_states(proprio=proprio) + + actions = rdt.predict_action( + lang_tokens=lang_embeddings, + lang_attn_mask=torch.ones( + lang_embeddings.shape[:2], dtype=torch.bool, device=DEVICE + ), + img_tokens=image_embeds, + state_tokens=states, # how can I get this? + action_mask=state_elem_mask.unsqueeze(1), # how can I get this? + ctrl_freqs=torch.tensor( + [25.0], device=DEVICE + ), # would this default work? + ) # (1, chunk_size, 128) + + # select the meaning action via STATE_INDICES + action = actions[ + :, :, state_indices + ] # (1, chunk_size, len(STATE_INDICES)) = (1, chunk_size, 7+ 1) + action = action.detach().float().to("cpu").numpy() + node.send_output("action", pa.array(action.ravel())) diff --git a/node-hub/dora-rdt-1b/pyproject.toml b/node-hub/dora-rdt-1b/pyproject.toml new file mode 100644 index 000000000..f4d3b0cde --- /dev/null +++ b/node-hub/dora-rdt-1b/pyproject.toml @@ -0,0 +1,48 @@ +[tool.poetry] +name = "dora-rdt-1b" +version = "0.3.8" +authors = ["Haixuan Xavier Tao "] +description = "Dora Node for RDT 1B" +readme = "README.md" + +packages = [{ include = "dora_rdt_1b" }] + +[tool.poetry.dependencies] +python = "^3.7" +dora-rs = "^0.3.6" +numpy = "< 2.0.0" +torch = "^2.4.0" +torchvision = "^0.19" +transformers = "^4.45" +qwen-vl-utils = "^0.0.2" +accelerate = "^0.33" +opencv-python = ">= 4.1.1" +modelscope = "^1.18.1" +packaging = "24.0" +wandb = "0.17.0" +diffusers = "0.27.2" +timm = "1.0.3" +sentencepiece = "0.2.0" +h5py = "3.11.0" +imgaug = "0.4.0" +huggingface_hub = "0.23.5" +# flash_attn = "^2.6.1" # Install using: pip install -U flash-attn --no-build-isolation + + +[tool.pylint.MASTER] +ignore-paths = '^dora_rdt_1b/RoboticsDiffusionTransformer.*$' + +[tool.poetry.dev-dependencies] +pytest = "^8.3.4" +pylint = "^3.3.2" + +[tool.black] +extend-exclude = 'dora_rdt_1b/RoboticsDiffusionTransformer' + + +[tool.poetry.scripts] +dora-rdt-1b = "dora_rdt_1b.main:main" + +[build-system] +requires = ["poetry-core>=1.8.0"] +build-backend = "poetry.core.masonry.api" diff --git a/node-hub/dora-rdt-1b/tests/conftest.py b/node-hub/dora-rdt-1b/tests/conftest.py new file mode 100644 index 000000000..46712c463 --- /dev/null +++ b/node-hub/dora-rdt-1b/tests/conftest.py @@ -0,0 +1,12 @@ +import pytest + + +def pytest_configure(): + pytest.rdt = None + pytest.lang_embeddings = None + pytest.image_processor = None + pytest.vision_encoder = None + pytest.image_embeds = None + pytest.state_elem_mask = None + pytest.states = None + pytest.STATE_INDICES = None diff --git a/node-hub/dora-rdt-1b/tests/test_dora_rdt_1b.py b/node-hub/dora-rdt-1b/tests/test_dora_rdt_1b.py new file mode 100644 index 000000000..b77e475f5 --- /dev/null +++ b/node-hub/dora-rdt-1b/tests/test_dora_rdt_1b.py @@ -0,0 +1,226 @@ +import pytest +import torch +import numpy as np +from PIL import Image +from torchvision import transforms +import os + + +CI = os.environ.get("CI") + +DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu" +DEVICE = os.getenv("DEVICE", DEVICE) +DEVICE = torch.device(DEVICE) +DTYPE = torch.float16 if DEVICE != "cpu" else torch.float32 + + +def test_import_main(): + # from dora_rdt_1b.main import main + + # Check that everything is working, and catch dora Runtime Exception as we're not running in a dora dataflow. + # with pytest.raises(RuntimeError): + # main() + import dora_rdt_1b.RoboticsDiffusionTransformer as _ + import dora_rdt_1b as _ + + +def test_download_policy(): + from dora_rdt_1b.main import get_policy + + ## in the future we should add this test within CI + if CI: + return + rdt = get_policy() + + pytest.rdt = rdt + + +def test_download_vision_model(): + from dora_rdt_1b.main import get_vision_model + + (vision_encoder, image_processor) = get_vision_model() + pytest.vision_encoder = vision_encoder + pytest.image_processor = image_processor + + +def test_download_language_embeddings(): + + ## in the future we should add this test within CI + if CI: + return + + from dora_rdt_1b.main import get_language_embeddings + + pytest.lang_embeddings = get_language_embeddings() + + +def test_load_dummy_image(): + + from dora_rdt_1b.main import config + + # Load pretrained model (in HF style) + image_processor = pytest.image_processor + vision_encoder = pytest.vision_encoder + + ## in the future we should add this test within CI + if CI: + return + + previous_image_path = "/path/to/img.jpeg" + # previous_image = None # if t = 0 + previous_image = Image.open(previous_image_path).convert("RGB") # if t > 0 + + current_image_path = "/path/to/img.jpeg" + current_image = Image.open(current_image_path).convert("RGB") + + # here I suppose you only have an image from exterior (e.g., 3rd person view) and you don't have any state information + # the images should arrange in sequence [exterior_image, right_wrist_image, left_wrist_image] * image_history_size (e.g., 2) + rgbs_lst = [[previous_image, None, None], [current_image, None, None]] + # if your have an right_wrist_image, then it should be + # rgbs_lst = [ + # [previous_image, previous_right_wrist_image, None], + # [current_image, current_right_wrist_image, None] + # ] + + # image pre-processing + # The background image used for padding + background_color = np.array( + [int(x * 255) for x in image_processor.image_mean], dtype=np.uint8 + ).reshape((1, 1, 3)) + background_image = ( + np.ones( + (image_processor.size["height"], image_processor.size["width"], 3), + dtype=np.uint8, + ) + * background_color + ) + + image_tensor_list = [] + for step in range(config["common"]["img_history_size"]): + rgbs = rgbs_lst[step % len(rgbs_lst)] + for rgb in rgbs: + if rgb is None: + # Replace it with the background image + image = Image.fromarray(background_image) + else: + image = rgb + + if config["dataset"].get("auto_adjust_image_brightness", False): + pixel_values = list(image.getdata()) + average_brightness = sum(sum(pixel) for pixel in pixel_values) / ( + len(pixel_values) * 255.0 * 3 + ) + if average_brightness <= 0.15: + image = transforms.ColorJitter(brightness=(1.75, 1.75))(image) + + if config["dataset"].get("image_aspect_ratio", "pad") == "pad": + + def expand2square(pil_img, background_color): + width, height = pil_img.size + if width == height: + return pil_img + elif width > height: + result = Image.new( + pil_img.mode, (width, width), background_color + ) + result.paste(pil_img, (0, (width - height) // 2)) + return result + else: + result = Image.new( + pil_img.mode, (height, height), background_color + ) + result.paste(pil_img, ((height - width) // 2, 0)) + return result + + image = expand2square( + image, tuple(int(x * 255) for x in image_processor.image_mean) + ) + image = image_processor.preprocess(image, return_tensors="pt")[ + "pixel_values" + ][0] + image_tensor_list.append(image) + + image_tensor = torch.stack(image_tensor_list, dim=0).to(DEVICE, dtype=DTYPE) + # encode images + image_embeds = vision_encoder(image_tensor).detach() + pytest.image_embeds = image_embeds.reshape( + -1, vision_encoder.hidden_size + ).unsqueeze(0) + + +def test_dummy_states(): + from dora_rdt_1b.main import config + + ## in the future we should add this test within CI + if CI: + return + + # suppose you do not have proprio + # it's kind of tricky, I strongly suggest adding proprio as input and further fine-tuning + B, N = 1, 1 # batch size and state history size + states = torch.zeros( + (B, N, config["model"]["state_token_dim"]), device=DEVICE, dtype=DTYPE + ) + + # if you have proprio, you can do like this + # format like this: [arm_joint_0_pos, arm_joint_1_pos, arm_joint_2_pos, arm_joint_3_pos, arm_joint_4_pos, arm_joint_5_pos, arm_joint_6_pos, gripper_open] + # proprio = torch.tensor([0, 1, 2, 3, 4, 5, 6, 0.5]).reshape((1, 1, -1)) + # states[:, :, STATE_INDICES] = proprio + + state_elem_mask = torch.zeros( + (B, config["model"]["state_token_dim"]), device=DEVICE, dtype=torch.bool + ) + from dora_rdt_1b.RoboticsDiffusionTransformer.configs.state_vec import ( + STATE_VEC_IDX_MAPPING, + ) + + # suppose you control in 7DOF joint position + STATE_INDICES = [ + STATE_VEC_IDX_MAPPING["arm_joint_0_pos"], + STATE_VEC_IDX_MAPPING["arm_joint_1_pos"], + STATE_VEC_IDX_MAPPING["arm_joint_2_pos"], + STATE_VEC_IDX_MAPPING["arm_joint_3_pos"], + STATE_VEC_IDX_MAPPING["arm_joint_4_pos"], + STATE_VEC_IDX_MAPPING["arm_joint_5_pos"], + STATE_VEC_IDX_MAPPING["arm_joint_6_pos"], + STATE_VEC_IDX_MAPPING["gripper_open"], + ] + + state_elem_mask[:, STATE_INDICES] = True + states, state_elem_mask = states.to(DEVICE, dtype=DTYPE), state_elem_mask.to( + DEVICE, dtype=DTYPE + ) + states = states[:, -1:, :] # only use the last state + pytest.states = states + pytest.state_elem_mask = state_elem_mask + pytest.STATE_INDICES = STATE_INDICES + + +def test_dummy_input(): + ## in the future we should add this test within CI + if CI: + return + + rdt = pytest.rdt + lang_embeddings = pytest.lang_embeddings + image_embeds = pytest.image_embeds + state_elem_mask = pytest.state_elem_mask + states = pytest.states + STATE_INDICES = pytest.STATE_INDICES + + actions = rdt.predict_action( + lang_tokens=lang_embeddings, + lang_attn_mask=torch.ones( + lang_embeddings.shape[:2], dtype=torch.bool, device=DEVICE + ), + img_tokens=image_embeds, + state_tokens=states, # how can I get this? + action_mask=state_elem_mask.unsqueeze(1), # how can I get this? + ctrl_freqs=torch.tensor([25.0], device=DEVICE), # would this default work? + ) # (1, chunk_size, 128) + + # select the meaning action via STATE_INDICES + action = actions[ + :, :, STATE_INDICES + ] # (1, chunk_size, len(STATE_INDICES)) = (1, chunk_size, 7+ 1) + print(action) diff --git a/node-hub/dora_rdt_1b/__init__.py b/node-hub/dora_rdt_1b/__init__.py new file mode 100644 index 000000000..ed4e21919 --- /dev/null +++ b/node-hub/dora_rdt_1b/__init__.py @@ -0,0 +1,19 @@ +import os +import sys +from pathlib import Path + +# Define the path to the README file relative to the package directory +readme_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "README.md") + +# Read the content of the README file +try: + with open(readme_path, "r", encoding="utf-8") as f: + __doc__ = f.read() +except FileNotFoundError: + __doc__ = "README file not found." + + +# Set up the import hook + +submodule_path = Path(__file__).resolve().parent / "RoboticsDiffusionTransformer" +sys.path.insert(0, str(submodule_path))