egogym 0.1.0__py3-none-any.whl

baselines/rum/utils/trajectory_vis.py

@@ -0,0 +1,128 @@
+import collections
+
+import einops
+import matplotlib.gridspec as gridspec
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+
+def generate_plots(ground_truths, predictions, sampled_images, to_plot=8, traj_index=0):
+    """
+    Generates plots comparing ground truth vs predictions and displays sampled images.
+
+    Parameters:
+    - ground_truths: A numpy array of shape (T, 7) containing ground truth actions.
+    - predictions: A numpy array of shape (T, 7) with model's predictions.
+    - sampled_images: A numpy array containing M sampled images.
+    """
+
+    # Initialize the figure
+    fig = plt.figure(figsize=(15, 15))
+    outer_grid = gridspec.GridSpec(8, 1, hspace=0.25, wspace=0.0)
+
+    # Top grid for images
+    M = sampled_images.shape[0]
+    chosen_indices = np.linspace(0, M - 1, to_plot, dtype=int)
+    chosen_images = sampled_images[chosen_indices]
+    top_grid = gridspec.GridSpecFromSubplotSpec(
+        1, len(chosen_images), subplot_spec=outer_grid[0, :], wspace=0.0
+    )
+
+    for i, img in enumerate(chosen_images):
+        ax = plt.Subplot(fig, top_grid[i])
+        ax.imshow(img)
+        ax.axis("off")
+        if i > 0:  # To further remove any potential whitespace
+            ax.set_yticklabels([])
+            ax.set_xticklabels([])
+            ax.set_xticks([])
+            ax.set_yticks([])
+        fig.add_subplot(ax)
+
+    # For each action dimension, plot ground truth vs prediction
+    dim = ground_truths.shape[1]
+    for i in range(dim):
+        ax = plt.Subplot(fig, outer_grid[i + 1, :])
+
+        gt_values = ground_truths[:, i]
+        pred_values = predictions[:, i]
+
+        ax.plot(gt_values, "g", label="Ground Truth")
+        ax.plot(pred_values, "r--", label="Prediction")
+
+        ax.legend(loc="upper right")
+        ax.set_title(f"Act Dim {i + 1}", fontsize=16)
+        fig.add_subplot(ax)
+
+    # save the plot
+    fig.savefig(f"trajectory_{traj_index}.png")
+
+
+@torch.no_grad()
+def visualize_trajectory(
+    model,
+    test_dataset,
+    device,
+    buffer_size=6,
+    n_visualized_trajectories=5,
+    goal_conditional=False,
+):
+    """
+    Visualizes the trajectory of the model on the test dataset.
+    M images are sampled from the test dataset and displayed.
+    T = end - start is the length of the trajectory.
+    start and end (index) indicate the window of the trajectory to visualize.
+    """
+    print(goal_conditional)
+    action_preds = []
+    ground_truth = []
+    images = []
+    image_buffer = collections.deque(maxlen=buffer_size)
+    test_dataset.set_include_trajectory_end(True)
+
+    i = 0
+    done_visualizing = 0
+    while (done_visualizing < n_visualized_trajectories) and (i < len(test_dataset)):
+        if goal_conditional:
+            (input_images, terminate), goals, *_, gt_actions = test_dataset[i]
+        else:
+            (input_images, terminate), *_, gt_actions = test_dataset[i]
+        input_images = input_images.float() / 255.0
+        image_buffer.append(input_images[-1])
+        img = input_images[-1]
+        images.append(einops.rearrange(img, "c h w -> h w c").cpu().detach().numpy())
+        ground_truth.append(gt_actions[-1])
+        if goal_conditional:
+            model_input = (
+                torch.stack(tuple(image_buffer), dim=0).unsqueeze(0).to(device),
+                torch.tensor(goals).unsqueeze(0).to(device),
+                torch.tensor(gt_actions).unsqueeze(0).to(device),
+            )
+        else:
+            model_input = (
+                torch.stack(tuple(image_buffer), dim=0).unsqueeze(0).to(device),
+                torch.tensor(gt_actions).unsqueeze(0).to(device),
+            )
+
+        out, _ = model.step(model_input)
+        action_preds.append(out.squeeze().cpu().detach().numpy())
+
+        if terminate:
+            action_preds = np.array(action_preds)
+            ground_truth = np.array(ground_truth)
+            images = np.array(images)
+
+            print(action_preds.shape, ground_truth.shape, images.shape)
+
+            generate_plots(
+                ground_truth, action_preds, images, traj_index=done_visualizing
+            )
+            done_visualizing += 1
+            # Reset everything.
+            action_preds = []
+            ground_truth = []
+            images = []
+            image_buffer = collections.deque(maxlen=buffer_size)
+
+        i += 1

baselines/rum/utils/zmq_utils.py

@@ -0,0 +1,281 @@
+import zmq
+import cv2
+import numpy as np
+import pickle
+import blosc as bl
+import threading
+import time
+from abc import ABC
+
+
+# ZMQ Sockets
+def create_push_socket(host, port):
+    context = zmq.Context()
+    socket = context.socket(zmq.PUSH)
+    socket.bind("tcp://{}:{}".format(host, port))
+    return socket
+
+
+def create_pull_socket(host, port):
+    context = zmq.Context()
+    socket = context.socket(zmq.PULL)
+    socket.setsockopt(zmq.CONFLATE, 1)
+    socket.bind("tcp://{}:{}".format(host, port))
+    return socket
+
+
+def create_response_socket(host, port):
+    content = zmq.Context()
+    socket = content.socket(zmq.REP)
+    socket.bind("tcp://{}:{}".format(host, port))
+    return socket
+
+
+def create_request_socket(host, port):
+    context = zmq.Context()
+    socket = context.socket(zmq.REQ)
+    socket.connect("tcp://{}:{}".format(host, port))
+    return socket
+
+
+# Pub/Sub classes for Keypoints
+class ZMQKeypointPublisher:
+    def __init__(self, host, port):
+        self._host, self._port = host, port
+        self._init_publisher()
+
+    def _init_publisher(self):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.PUB)
+        self.socket.bind("tcp://{}:{}".format(self._host, self._port))
+
+    def pub_keypoints(self, keypoint_array, topic_name):
+        """
+        Process the keypoints into a byte stream and input them in this function
+        """
+        buffer = pickle.dumps(keypoint_array, protocol=-1)
+        self.socket.send(bytes("{} ".format(topic_name), "utf-8") + buffer)
+
+    def stop(self):
+        print("Closing the publisher socket in {}:{}.".format(self._host, self._port))
+        self.socket.close()
+        self.context.term()
+
+
+# Keypoint Subscriber
+class ZMQKeypointSubscriber(threading.Thread):
+    def __init__(self, host, port, topic):
+        self._host, self._port, self._topic = host, port, topic
+        self._init_subscriber()
+
+        # Topic chars to remove
+        self.strip_value = bytes("{} ".format(self._topic), "utf-8")
+
+    def _init_subscriber(self):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.SUB)
+        self.socket.setsockopt(zmq.CONFLATE, 1)
+        self.socket.connect("tcp://{}:{}".format(self._host, self._port))
+        self.socket.setsockopt(zmq.SUBSCRIBE, bytes(self._topic, "utf-8"))
+
+    def recv_keypoints(self, flags=None):
+        if flags is None:
+            raw_data = self.socket.recv()
+            raw_array = raw_data.lstrip(self.strip_value)
+            return pickle.loads(raw_array)
+        else:  # For possible usage of no blocking zmq subscriber
+            try:
+                raw_data = self.socket.recv(flags)
+                raw_array = raw_data.lstrip(self.strip_value)
+                return pickle.loads(raw_array)
+            except zmq.Again:
+                return None
+
+    def stop(self):
+        print("Closing the subscriber socket in {}:{}.".format(self._host, self._port))
+        self.socket.close()
+        self.context.term()
+
+
+# Pub/Sub classes for storing data from Realsense Cameras
+class ZMQCameraPublisher:
+    def __init__(self, host, port):
+        self._host, self._port = host, port
+        self._init_publisher()
+
+    def _init_publisher(self):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.PUB)
+        print("tcp://{}:{}".format(self._host, self._port))
+        self.socket.bind("tcp://{}:{}".format(self._host, self._port))
+
+    def pub_intrinsics(self, array):
+        self.socket.send(b"intrinsics " + pickle.dumps(array, protocol=-1))
+
+    def pub_rgb_image(self, rgb_image, timestamp):
+        _, buffer = cv2.imencode(".jpg", rgb_image, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
+        data = dict(timestamp=timestamp, rgb_image=buffer.tobytes())
+        self.socket.send(b"rgb_image " + pickle.dumps(data, protocol=-1))
+
+    def pub_depth_image(self, depth_image, timestamp):
+        compressed_depth = bl.pack_array(
+            depth_image, cname="zstd", clevel=1, shuffle=bl.NOSHUFFLE
+        )
+        data = dict(timestamp=timestamp, depth_image=compressed_depth)
+        self.socket.send(b"depth_image " + pickle.dumps(data, protocol=-1))
+
+    def pub_image_and_depth(self, rgb_image, depth_image, timestamp):
+        _, buffer = cv2.imencode(".jpg", rgb_image, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
+        compressed_depth = bl.pack_array(
+            depth_image, cname="zstd", clevel=1, shuffle=bl.NOSHUFFLE
+        )
+        data = dict(
+            timestamp=timestamp,
+            rgb_image=buffer.tobytes(),
+            depth_image=compressed_depth,
+        )
+        self.socket.send(b"image_and_depth " + pickle.dumps(data, protocol=-1))
+
+    def stop(self):
+        print("Closing the publisher socket in {}:{}.".format(self._host, self._port))
+        self.socket.close()
+        self.context.term()
+
+
+class ZMQCameraSubscriber:
+    def __init__(self, host, port, topic_type):
+        self._host, self._port, self._topic_type = host, port, topic_type
+        self._init_subscriber()
+
+    def _init_subscriber(self):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.SUB)
+        self.socket.setsockopt(zmq.CONFLATE, 1)
+        self.socket.connect("tcp://{}:{}".format(self._host, self._port))
+
+        if self._topic_type == "Intrinsics":
+            self.socket.setsockopt(zmq.SUBSCRIBE, b"intrinsics")
+        elif self._topic_type == "RGB":
+            self.socket.setsockopt(zmq.SUBSCRIBE, b"rgb_image")
+        elif self._topic_type == "Depth":
+            self.socket.setsockopt(zmq.SUBSCRIBE, b"depth_image")
+        elif self._topic_type == "RGBD":
+            self.socket.setsockopt(zmq.SUBSCRIBE, b"image_and_depth")
+
+    def recv_intrinsics(self):
+        raw_data = self.socket.recv()
+        raw_array = raw_data.lstrip(b"intrinsics ")
+        return pickle.loads(raw_array)
+
+    def recv_rgb_image(self):
+        raw_data = self.socket.recv()
+        data = raw_data.lstrip(b"rgb_image ")
+        data = pickle.loads(data)
+        encoded_data = np.frombuffer(data["rgb_image"], np.uint8)
+        return cv2.imdecode(encoded_data, 1), data["timestamp"]
+
+    def recv_depth_image(self):
+        raw_data = self.socket.recv()
+        striped_data = raw_data.lstrip(b"depth_image ")
+        data = pickle.loads(striped_data)
+        depth_image = bl.unpack_array(data["depth_image"])
+        return np.array(depth_image, dtype=np.float32), data["timestamp"]
+
+    def recv_image_and_depth(self):
+        raw_data = self.socket.recv()
+        striped_data = raw_data.lstrip(b"image_and_depth ")
+        data = pickle.loads(striped_data)
+        encoded_data = np.frombuffer(data["rgb_image"], np.uint8)
+        rgb_image = cv2.imdecode(encoded_data, 1)
+        depth_image = bl.unpack_array(data["depth_image"])
+        return rgb_image, np.array(depth_image, dtype=np.float32), data["timestamp"]
+
+    def stop(self):
+        print("Closing the subscriber socket in {}:{}.".format(self._host, self._port))
+        self.socket.close()
+        self.context.term()
+
+
+# Publisher for image visualizers
+class ZMQCompressedImageTransmitter(object):
+    def __init__(self, host, port):
+        self._host, self._port = host, port
+        # self._init_push_socket()
+        self._init_publisher()
+
+    def _init_publisher(self):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.PUB)
+        self.socket.bind("tcp://{}:{}".format(self._host, self._port))
+
+    def _init_push_socket(self):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.PUSH)
+        self.socket.bind("tcp://{}:{}".format(self._host, self._port))
+
+    def send_image(self, rgb_image):
+        _, buffer = cv2.imencode(".jpg", rgb_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50])
+        self.socket.send(buffer.tobytes())
+
+    def stop(self):
+        print("Closing the publisher in {}:{}.".format(self._host, self._port))
+        self.socket.close()
+        self.context.term()
+
+
+class ZMQCompressedImageReciever(threading.Thread):
+    def __init__(self, host, port):
+        self._host, self._port = host, port
+        # self._init_pull_socket()
+        self._init_subscriber()
+
+    def _init_subscriber(self):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.SUB)
+        self.socket.setsockopt(zmq.CONFLATE, 1)
+        self.socket.connect("tcp://{}:{}".format(self._host, self._port))
+        self.socket.subscribe("")
+
+    def _init_pull_socket(self):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.PULL)
+        self.socket.setsockopt(zmq.CONFLATE, 1)
+        self.socket.connect("tcp://{}:{}".format(self._host, self._port))
+
+    def recv_image(self):
+        raw_data = self.socket.recv()
+        encoded_data = np.frombuffer(raw_data, np.uint8)
+        decoded_frame = cv2.imdecode(encoded_data, 1)
+        return decoded_frame
+
+    def stop(self):
+        print("Closing the subscriber socket in {}:{}.".format(self._host, self._port))
+        self.socket.close()
+        self.context.term()
+
+
+class FrequencyTimer:
+    FREQ_1KHZ = 1e3
+
+    def __init__(self, frequency_rate):
+        self.time_available = 1e9 / frequency_rate
+
+    def start_loop(self):
+        self.start_time = time.time_ns()
+
+    def end_loop(self):
+        wait_time = self.time_available + self.start_time
+
+        while time.time_ns() < wait_time:
+            time.sleep(1 / FrequencyTimer.FREQ_1KHZ)
+
+
+class ProcessInstantiator(ABC):
+    def __init__(self):
+        self.processes = []
+
+    def _start_component(self, configs):
+        raise NotImplementedError("Function not implemented!")
+
+    def get_processes(self):
+        return self.processes

baselines/rum_policy.py

@@ -0,0 +1,108 @@
+import cv2
+import torch
+import torch.nn.functional as F
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+from egogym.policies.base_policy import BasePolicy
+from baselines.rum.policy import Policy
+
+class RUMPolicy(BasePolicy):
+    
+    def __init__(self, config=None, policy=None, **kwargs):
+        super().__init__()
+
+        if config is None:
+            config = get_config()
+        
+        for key, value in kwargs.items():
+            config[key] = value
+            
+        self.config = config
+        self.name = "rum"
+        self.device = config.get("device", "cuda" if torch.cuda.is_available() else "cpu")
+        if policy is not None:
+            self.policy = policy
+        else:
+            self.policy = Policy(model_path=config["model_path"], device=self.device)
+        self.gripper_threshold = config["gripper_threshold"]
+        self.grasping_style = config["grasping_style"]
+        self.num_envs = None
+        self.grasped = None
+        self.desired_axis_t = torch.tensor([
+            [-1,  0,  0],
+            [ 0,  0, -1],
+            [ 0, -1,  0]
+        ], dtype=torch.float32, device=self.device)
+        self.grasp_override_t = torch.tensor([0.00, 0.18, 0.04], dtype=torch.float32, device=self.device)
+
+    def get_action(self, obs):
+
+        if self.num_envs is None:
+            if len(obs["rgb_ego"].shape) == 4:
+                self.num_envs = obs["rgb_ego"].shape[0]
+            else:
+                self.num_envs = 1
+            self.grasped = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
+        
+        if self.num_envs == 1:
+            obs = {key: obs[key][np.newaxis, ...] if isinstance(obs[key], np.ndarray) else obs[key] for key in obs}
+
+        T_world_camera = torch.from_numpy(obs["camera_pose"]).view(self.num_envs, 4, 4).to(self.device)
+        if "object_pose" in obs:
+            T_world_object = torch.from_numpy(obs["object_pose"]).view(self.num_envs, 4, 4).to(self.device)
+        else:
+            T_world_object = torch.from_numpy(obs["handle_pose"]).view(self.num_envs, 4, 4).to(self.device)
+        
+        R_world_camera = T_world_camera[:, :3, :3]
+        t_world_camera = T_world_camera[:, :3, 3]
+        p_world = T_world_object[:, :3, 3]
+        R_camera_world = R_world_camera.transpose(1, 2)
+        p_camera = torch.bmm(R_camera_world, (p_world - t_world_camera).unsqueeze(-1)).squeeze(-1)
+        object_3d_position = p_camera @ self.desired_axis_t.T
+        object_3d_position[self.grasped] = self.grasp_override_t
+        
+        rgb_tensor = torch.from_numpy(obs["rgb_ego"]).permute(0, 3, 1, 2).to(self.device).float()
+        rgb_resized = F.interpolate(rgb_tensor, size=(224, 224), mode='bilinear', align_corners=False)
+        rgb_ego = rgb_resized.permute(0, 2, 3, 1).byte().cpu().numpy()
+
+        rum_obs = {
+            "rgb_ego": rgb_ego,
+            "object_3d_position": object_3d_position.cpu().numpy(),
+        }
+        action_tensors = self.policy.infer(rum_obs)
+        
+        action_tensors_t = torch.from_numpy(action_tensors).to(self.device)
+        R_action = R.from_euler("xyz", action_tensors[:, 3:6]).as_matrix()
+        R_action_t = torch.from_numpy(R_action).to(self.device, dtype=torch.float32)
+        
+        action_pose = torch.eye(4, device=self.device).unsqueeze(0).repeat(self.num_envs, 1, 1)
+        action_pose[:, :3, :3] = R_action_t
+        action_pose[:, :3, 3] = action_tensors_t[:, :3]
+        action_pose[:, :3, :3] = self.desired_axis_t.T @ action_pose[:, :3, :3] @ self.desired_axis_t
+        action_pose[:, :3, 3] = action_pose[:, :3, 3] @ self.desired_axis_t
+        
+        self.grasped = torch.maximum(self.grasped, action_tensors_t[:, 6] < self.gripper_threshold)
+        grasp_action = self.grasped.unsqueeze(1).float()
+        if self.grasping_style == "continuous":
+            grasp_action = 1 - action_tensors_t[:, 6].unsqueeze(1)
+
+        actions = torch.cat([action_pose.view(self.num_envs, -1), grasp_action], dim=1)
+
+        return actions.cpu().numpy()
+    
+    def reset(self, indicies=None):
+        self.policy.reset(indicies)
+        if indicies is not None and self.num_envs is not None:
+            self.grasped[indicies] = False
+        elif self.num_envs is not None:
+            self.grasped = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
+
+
+def get_config():
+    config = {
+        "model_path": "checkpoints/rum_pick.pt",
+        "gripper_threshold": 0.7,
+        "grasping_style": "binary",
+        "device": None,
+    }
+    return config

egogym/__init__.py

@@ -0,0 +1,8 @@
+import gymnasium as gym
+from .tasks.pick import PickTask
+from .tasks.open import OpenTask
+from .tasks.close import CloseTask
+
+gym.register(id="Egogym-Pick-v0", entry_point=PickTask)
+gym.register(id="Egogym-Open-v0", entry_point=OpenTask)
+gym.register(id="Egogym-Close-v0", entry_point=CloseTask)