robits 0.5.0__py3-none-any.whl

robits/agents/base_agent.py

@@ -0,0 +1,63 @@
+from abc import ABC
+from abc import abstractmethod
+
+import logging
+
+import numpy as np
+
+
+logger = logging.getLogger(__name__)
+
+
+class BaseAgent(ABC):
+    """ """
+
+    agent_name: str
+
+    lang_goal: str
+
+    def prepare_observation(self, obs, i, episode_length):
+
+        from clip import tokenize
+        import torch
+
+        logger.info("Using language goal %s", self.lang_goal)
+
+        obs["lang_goal_tokens"] = tokenize([self.lang_goal])[0].numpy()
+
+        elapsed_time = (1.0 - (i / float(episode_length - 1))) * 2.0 - 1.0
+        gripper_joint_positions = obs["gripper_joint_positions"][
+            0:1
+        ]  # gripper is normalized here
+        logger.info("timestamp: %.2f", elapsed_time)
+
+        obs["ignore_collisions"] = np.array([0])
+        obs["low_dim_state"] = np.concatenate(
+            [
+                gripper_joint_positions,
+                gripper_joint_positions,
+                gripper_joint_positions,
+                elapsed_time,
+            ],
+            axis=None,
+        )
+
+        for k, v in obs.items():
+            if v is None:
+                logger.info("No values for key %s", k)
+                # obs[k] = torch.tensor([np.zeros(4)], device=self.device).unsqueeze(0)
+                continue
+
+            if isinstance(v, np.ndarray):
+                v = v.astype(np.float32)
+                logger.debug("Item %s has shape %s", k, v.shape)
+            else:
+                logger.debug("Key %s is not a numpy array", k)
+
+            obs[k] = torch.tensor([v], device=self.device).unsqueeze(0)
+
+        return obs
+
+    @abstractmethod
+    def get_action(self, step, observation):
+        pass

robits/agents/mello_agent.py

@@ -0,0 +1,72 @@
+from typing import Optional
+import logging
+import time
+import requests
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class MelloConfig:
+
+    max_delta: float = 0.25
+
+    initial_joint_positions: Optional[np.ndarray] = None
+
+
+class MelloAgent:
+
+    def __init__(self, device_addr, robot):
+        self.device_addr = device_addr
+        self.robot = robot
+        self.max_delta = 0.25
+
+    def get_mello_data(self):
+        data = requests.get(self.device_addr).json()
+        joint_positions = np.array(data["joint_positions"])
+        joint_positions = np.deg2rad(joint_positions)
+        joint_positions *= -1
+
+        joint_velocities = np.array(data["joint_velocities"])
+        joint_velocities = np.deg2rad(joint_velocities)
+        joint_velocities *= -1
+        return joint_positions, joint_velocities
+
+    def get_mello_joint_positions(self):
+        joint_positions = np.array(
+            requests.get(self.device_addr).json()["joint_positions"]
+        )
+        joint_positions = np.deg2rad(joint_positions)
+        joint_positions *= -1
+        return joint_positions
+
+    def get_action(self) -> np.ndarray:
+        mello_joint_positions = self.get_mello_joint_positions()
+        robot_joint_positions = self.robot.get_proprioception_data(False, False)[
+            "joint_positions"
+        ]
+        delta = mello_joint_positions - robot_joint_positions
+        delta = np.clip(delta, -self.max_delta, self.max_delta)
+        new_joint_positions = robot_joint_positions + delta
+        return new_joint_positions
+
+    def wait_for_pose(self):
+        delta = np.ones_like(self.get_mello_joint_positions()) * self.max_delta
+        while np.any(delta >= self.max_delta):
+            mello_joint_positions = self.get_mello_joint_positions()
+            robot_joint_positions = self.robot.get_joint_obs(False, False)[
+                "joint_positions"
+            ]
+            delta = np.abs(mello_joint_positions - robot_joint_positions)
+            logger.info(
+                "Delta between robot and mello is too large %s. Should be with %s",
+                np.rad2deg(delta),
+                np.rad2deg(self.max_delta),
+            )
+            time.sleep(0.25)
+        logger.info("Done.")

robits/agents/sam2act_agent.py

@@ -0,0 +1,29 @@
+import logging
+
+import torch
+
+from sam2act.eval import load_agent
+from robits.agents.base_agent import BaseAgent
+
+logger = logging.getLogger(__name__)
+
+
+class SAM2Act(BaseAgent):
+
+    def __init__(self, model_path):
+        self.device = torch.device("cuda:0")
+
+        self.agent = load_agent(model_path)
+        self.agent.load_clip()
+        self.agent.eval()
+        self.lang_goal = "push the buttons in the following order: red, green, blue"
+
+    def get_action(self, step, observation):
+
+        if not observation:
+            logger.error("Nothing todo.")
+            return
+
+        with torch.jit.optimized_execution(False):
+            act_result = self.agent.act(step, observation, deterministic=True)
+        return act_result

robits/app/__init__.py

@@ -0,0 +1,13 @@
+"""
+Interactive applications for the RoBits framework.
+
+This package provides user interface applications for interacting with
+robots, cameras, and other components of the RoBits system. These applications
+use terminal-based UI libraries to provide graphical interfaces that work
+in both terminal and GUI environments.
+
+Components:
+
+- robot_app: An interactive TUI for controlling robot arms
+- extrinsics_app: An application for calibrating camera extrinsics
+"""

robits/app/extrinsics_app.py

@@ -0,0 +1,339 @@
+#!/usr/bin/env python3
+
+"""
+App for extrinsic calibration
+"""
+import sys
+import os
+
+import time
+from threading import Thread
+
+from datetime import datetime
+
+
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+import open3d as o3d
+
+from textual.app import App
+from textual.app import ComposeResult
+from textual.containers import Vertical
+from textual.containers import Horizontal
+from textual.widgets import Label
+from textual.widgets import Button
+from textual_slider import Slider
+from textual.widgets import Select
+
+
+from robits.core.config_manager import config_manager
+from robits.core.factory import CameraFactory
+from robits.core.config import CameraCalibration
+from robits.utils import vision_utils
+
+
+class SimplePCDViewer:
+
+    def __init__(self) -> None:
+        self.viewer_running = True
+        self.prev_transform = np.identity(4)
+        self.table = None
+
+        self.point_cloud = o3d.geometry.PointCloud()
+        # self.point_cloud.points = o3d.utility.Vector3dVector(np.zeros((10, 3)))
+        self.prev_transform = np.identity(4)
+
+    def apply_transformation(self, transformation: np.ndarray) -> None:
+        """
+        Apply a transformation matrix to the current point cloud.
+
+        :param transformation: the transformation to apply
+        """
+        self.point_cloud.transform(np.linalg.inv(self.prev_transform))
+        self.point_cloud.transform(transformation)
+        self.prev_transform = transformation
+
+    def update_point_cloud(self, camera) -> None:
+        """
+        Get a new point clodu from the camera
+        :param camera: the camera to retrieve the point cloud from
+        """
+        camera_data, _ = camera.get_camera_data()
+        pcd = vision_utils.depth_to_pcd(camera_data, camera)
+        self.point_cloud.points = pcd.points
+        self.point_cloud.colors = pcd.colors
+
+    def show_table(self):
+        pass
+
+    def get_table(self):
+        table = o3d.geometry.TriangleMesh.create_box(width=0.6, height=1.5, depth=0.01)
+        table = table.translate(np.array([0.3, 0.0 - 1.5 / 2, 0.295 - 0.01 / 2.0]))
+        return table
+
+    def run(self) -> None:
+        """
+        Start the Open3D viewer in a separate thread.
+
+        .. todo:: this breaks compatibility to macOS as the gui is supposed to run in the main thread.
+        """
+        vis = o3d.visualization.Visualizer()
+        vis.create_window(window_name="Point Cloud Viewer")
+
+        coordinate_system = o3d.geometry.TriangleMesh.create_coordinate_frame(
+            size=0.4, origin=[0, 0, 0]
+        )
+        vis.add_geometry(coordinate_system)
+        vis.add_geometry(self.get_table())
+        vis.add_geometry(self.point_cloud)
+
+        vis.add_geometry(
+            o3d.geometry.TriangleMesh.create_box(
+                width=1.0, height=0.001, depth=1.0
+            ).paint_uniform_color([1.0, 0.0, 0.0])
+        )
+        vis.add_geometry(
+            o3d.geometry.TriangleMesh.create_box(
+                width=1.0, height=1.0, depth=0.001
+            ).paint_uniform_color([0.0, 1.0, 0.0])
+        )
+        vis.add_geometry(
+            o3d.geometry.TriangleMesh.create_box(width=0.001, height=1.0, depth=1.0)
+            .paint_uniform_color([0.0, 0.0, 1.0])
+            .translate(np.array([0.4687, 0.0, 0.0]))
+        )
+
+        while self.viewer_running:
+            vis.update_geometry(self.point_cloud)
+            vis.poll_events()
+            vis.update_renderer()
+            time.sleep(0.05)
+
+        vis.destroy_window()
+
+    def stop(self) -> None:
+        """
+        Stop the Open3D viewer.
+        """
+        self.viewer_running = False
+
+
+class CameraCalibrationApp(App):
+
+    def __init__(self, viewer: SimplePCDViewer):
+        super().__init__()
+        self.viewer = viewer
+
+    def compose(self) -> ComposeResult:
+        yield Vertical(
+            Label("current transform", id="label_transform"),
+            Horizontal(
+                Label("Select Camera"),
+                Select.from_values(
+                    values=config_manager.available_cameras, id="camera_select"
+                ),
+                Button(label="Connect", id="connect"),
+            ),
+            Label("Camera Calibration - Translation Sliders"),
+            Horizontal(
+                Vertical(
+                    Label("Translation X"),
+                    Slider(min=-3.0, max=3.0, value=0.0, step=0.0025, id="trans_x"),
+                    Label(id="label_x"),
+                ),
+                Vertical(
+                    Label("Translation Y"),
+                    Slider(min=-3.0, max=3.0, value=0.0, step=0.0025, id="trans_y"),
+                    Label(id="label_y"),
+                ),
+                Vertical(
+                    Label("Translation Z"),
+                    Slider(min=-3.0, max=3.0, value=0.0, step=0.0025, id="trans_z"),
+                    Label(id="label_z"),
+                ),
+            ),
+            Label("Camera Calibration - Rotation Sliders"),
+            Horizontal(
+                Vertical(
+                    Label("Rotation Roll"),
+                    Slider(min=-180, max=180, value=0, step=1, id="rot_roll"),
+                    Label(id="label_roll"),
+                ),
+                Vertical(
+                    Label("Rotation Pitch"),
+                    Slider(min=-180, max=180, value=0, step=1, id="rot_pitch"),
+                    Label(id="label_pitch"),
+                ),
+                Vertical(
+                    Label("Rotation Yaw"),
+                    Slider(min=-180, max=180, value=0, step=1, id="rot_yaw"),
+                    Label(id="label_yaw"),
+                ),
+            ),
+            Label("Actions"),
+            Horizontal(
+                Button(label="Save", id="save"),
+                Button(label="Refine", id="refine"),
+                Button(label="Exit", id="exit"),
+            ),
+        )
+
+    def on_mount(self):
+        # Start the Open3D viewer in a separate thread. Does not work for macOS
+        Thread(target=self.viewer.run, daemon=True).start()
+
+    def on_button_pressed(self, message: Button.Pressed) -> None:
+        if message.button.id == "exit":
+            sys.exit(0)
+        elif message.button.id == "connect":
+            selected_camera = self.query_one("#camera_select", Select).value
+            if selected_camera:
+                camera = CameraFactory(selected_camera).build()
+                self.viewer.update_point_cloud(camera)
+                transformation = np.linalg.inv(camera.calibration.extrinsics)
+                self.update_sliders(transformation)
+                self.viewer.apply_transformation(self.get_transformation())
+        elif message.button.id == "save":
+            selected_camera = self.query_one("#camera_select", Select).value
+            self.save_camera_calibration(selected_camera)
+        elif message.button.id == "refine":
+            self.estimate_transformation()
+
+    def update_sliders(self, extrinsics):
+        rpy = R.from_matrix(extrinsics[:3, :3]).as_euler("ZYX", degrees=True)
+        self.query_one("#trans_x", Slider).value = extrinsics[0, 3]
+        self.query_one("#trans_y", Slider).value = extrinsics[1, 3]
+        self.query_one("#trans_z", Slider).value = extrinsics[2, 3]
+        self.query_one("#rot_roll", Slider).value = rpy[2]
+        self.query_one("#rot_pitch", Slider).value = rpy[1]
+        self.query_one("#rot_yaw", Slider).value = rpy[0]
+
+    def draw_registration_result(self, source, target, transformation):
+        import copy
+
+        source_temp = copy.deepcopy(source)
+        target_temp = copy.deepcopy(target)
+        source_temp.paint_uniform_color([1, 0.706, 0])
+        target_temp.paint_uniform_color([0, 0.651, 0.929])
+        # source_temp.transform(transformation)
+        coordinate_system = o3d.geometry.TriangleMesh.create_coordinate_frame(
+            size=0.4, origin=[0, 0, 0]
+        )
+        o3d.visualization.draw_geometries([source_temp, target_temp, coordinate_system])
+
+    def estimate_transformation(self):
+        """
+        .. todo:: not implemented yet
+        """
+
+        pcd = self.viewer.point_cloud
+        """
+        if False:
+            plane_model, inliers = pcd.segment_plane(
+                distance_threshold=0.03, ransac_n=5000, num_iterations=1000
+            )
+            inlier_cloud = pcd.select_by_index(inliers)
+            inlier_cloud.paint_uniform_color([1.0, 0, 0])
+            outlier_cloud = pcd.select_by_index(inliers, invert=True)
+            extracted_table_pcd = inlier_cloud
+        
+        """
+
+        extracted_table_pcd = pcd
+
+        table_height = 0.25
+        x = np.linspace(0, 0.5, 100)
+        y = np.linspace(-0.7, 0.7, 100)
+        mesh_x, mesh_y = np.meshgrid(x, y)
+        xyz = np.zeros((100 * 100, 3))
+        xyz[:, 0] = mesh_x.reshape(-1)
+        xyz[:, 1] = mesh_y.reshape(-1)
+        xyz[:, 2] = table_height
+
+        target = o3d.geometry.PointCloud()
+        target.points = o3d.utility.Vector3dVector(xyz)
+        target.paint_uniform_color([1.0, 0.0, 0.0])
+
+        trans_init = self.get_transformation()
+        threshold = 0.02
+        self.draw_registration_result(extracted_table_pcd, target, trans_init)
+        reg_p2p = o3d.pipelines.registration.registration_icp(
+            extracted_table_pcd,
+            target,
+            threshold,
+            trans_init,
+            o3d.pipelines.registration.TransformationEstimationPointToPoint(),
+        )
+        # self.update_sliders(reg_p2p.transformation)
+        self.draw_registration_result(
+            extracted_table_pcd, target, reg_p2p.transformation
+        )
+
+    def get_transformation(self) -> np.ndarray:
+        trans_x = self.query_one("#trans_x", Slider).value
+        trans_y = self.query_one("#trans_y", Slider).value
+        trans_z = self.query_one("#trans_z", Slider).value
+        roll = self.query_one("#rot_roll", Slider).value
+        pitch = self.query_one("#rot_pitch", Slider).value
+        yaw = self.query_one("#rot_yaw", Slider).value
+
+        self.query_one("#label_x", Label).update(content=f"{trans_x:.2f}")
+        self.query_one("#label_y", Label).update(content=f"{trans_y:.2f}")
+        self.query_one("#label_z", Label).update(content=f"{trans_z:.2f}")
+
+        self.query_one("#label_roll", Label).update(content=f"{roll:.2f}")
+        self.query_one("#label_pitch", Label).update(content=f"{pitch:.2f}")
+        self.query_one("#label_yaw", Label).update(content=f"{yaw:.2f}")
+
+        # Create transformation matrix
+        transformation = np.identity(4)
+        transformation[:3, 3] = [trans_x, trans_y, trans_z]
+        transformation[:3, :3] = R.from_euler(
+            "ZYX", [yaw, pitch, roll], degrees=True
+        ).as_matrix()
+        # transformation = np.linalg.inv(transformation)
+
+        self.query_one("#label_transform", Label).update(content=f"{transformation}")
+
+        return transformation
+
+    def on_slider_changed(self, message: Slider.Changed) -> None:
+        transformation = self.get_transformation()
+
+        # Apply transformation to the point cloud
+        self.viewer.apply_transformation(transformation)
+
+    def on_shutdown(self):
+        # Stop the Open3D viewer when the app is closed
+        self.viewer.stop()
+
+    def save_camera_calibration(self, selected_camera):
+
+        if not selected_camera:
+            return
+
+        user_config_dir = config_manager.get_user_config_dir()
+        extrinsics = np.linalg.inv(self.get_transformation())
+
+        camera = CameraFactory(selected_camera).build()
+        config_dict = camera.calibration.to_dict()
+        config_dict["extrinsics"] = extrinsics.tolist()
+        config_dict["date_updated"] = datetime.now().isoformat()
+        new_calibration = CameraCalibration(**config_dict)
+        config_path = os.path.join(
+            user_config_dir, f"calibration_{camera.camera_name}_camera.json"
+        )
+        new_calibration.save_config(config_path)
+
+
+if __name__ == "__main__":
+    # import asyncio
+    # loop = asyncio.new_event_loop()
+    # asyncio.set_event_loop(loop)
+
+    print("Viewer initialized")
+    viewer = SimplePCDViewer()
+    print("Creating app")
+    app = CameraCalibrationApp(viewer)
+    print("Running app")
+    app.run()

robits/app/misc/slider_widget.py

@@ -0,0 +1,88 @@
+from textual.widget import Widget
+from textual.reactive import reactive
+from textual.events import Key
+from textual.message import Message
+from textual.containers import Horizontal
+from textual.containers import Vertical
+
+from textual.widgets import Button, Static
+from textual.app import ComposeResult
+
+
+class Slider(Widget):
+    """A simple horizontal slider widget with + and - buttons."""
+
+    value = reactive(0.0)
+
+    def __init__(self, min=0.0, max=1.0, value=0.0, step=0.1, id=None, name=None):
+        self._mounted = False
+        super().__init__(id=id, name=name)
+        self.min = min
+        self.max = max
+        self.step = step
+        self.bar_width = 30
+        self.value = value
+
+    def compose(self) -> ComposeResult:
+        decrease_btn = Button("-", id="decrease")
+        increase_btn = Button("+", id="increase")
+
+        for btn in (decrease_btn, increase_btn):
+            btn.styles.width = 3
+            btn.styles.height = 3
+            btn.styles.min_width = 3
+            btn.styles.min_height = 3
+            btn.styles.max_width = 3
+            btn.styles.max_height = 3
+
+        # Bar and label
+        bar = Static(self._render_bar(), id="bar", expand=True)
+        value_label = Static(f"{self.value:.3f}", id="value_label", expand=True)
+        bar.styles.text_align = "center"
+        value_label.styles.text_align = "center"
+
+        center_column = Vertical(value_label, bar, id="slider-column")
+        center_column.styles.align_horizontal = "center"
+
+        yield Horizontal(decrease_btn, center_column, increase_btn, id="slider-row")
+
+    def on_mount(self) -> None:
+        self._mounted = True
+        self.styles.width = self.bar_width + 10
+
+    def _render_bar(self) -> str:
+        normalized = (self.value - self.min) / (self.max - self.min)
+        filled_chars = int(normalized * self.bar_width)
+        empty_chars = self.bar_width - filled_chars
+        return f"[{'█' * filled_chars}{' ' * empty_chars}]"
+
+    def on_button_pressed(self, event: Button.Pressed) -> None:
+        if event.button.id == "decrease":
+            self.decrease()
+        elif event.button.id == "increase":
+            self.increase()
+
+    def decrease(self):
+        self.value = max(self.min, self.value - self.step)
+        self.post_message(Slider.Changed(self))
+
+    def increase(self):
+        self.value = min(self.max, self.value + self.step)
+        self.post_message(Slider.Changed(self))
+
+    def on_key(self, event: Key) -> None:
+        if event.key == "left":
+            self.decrease()
+        elif event.key == "right":
+            self.increase()
+
+    def watch_value(self, old: float, new: float) -> None:
+        if self._mounted:
+            self.query_one("#bar", Static).update(self._render_bar())
+            self.query_one("#value_label", Static).update(f"{new:.2f}")
+
+    class Changed(Message):
+        """Message sent when the slider value changes."""
+
+        def __init__(self, sender: "Slider"):
+            super().__init__()