egogym 0.1.0__py3-none-any.whl

baselines/pi_policy.py

@@ -0,0 +1,110 @@
+import numpy as np
+import time
+from egogym.policies.base_policy import BasePolicy
+from egogym.utils import resize_with_pad
+
+class PIPolicy(BasePolicy):
+    
+    def __init__(self, config=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 = "pi"
+        self.host = self.config["host"]
+        self.port = self.config["port"]
+        self.grasping_style = self.config["grasping_style"]
+        self.buffer_length = self.config["buffer_length"]
+        self.batch_size = 0
+
+        try:
+            from openpi_client import websocket_client_policy
+        except ImportError:
+            raise ImportError("Please install the openpi-client.")
+        
+        max_retries = 5
+        for attempt in range(max_retries):
+            try:
+                self.model = websocket_client_policy.WebsocketClientPolicy(
+                    host=self.host,
+                    port=self.port,
+                )
+                print(f"Connected to PI model at {self.host}:{self.port}")
+                break
+            except Exception as e:
+                if attempt < max_retries - 1:
+                    print(f"Connection attempt {attempt + 1} failed: {e}. Retrying...")
+                    time.sleep(1)
+                else:
+                    print(f"Failed to connect to remote model after {max_retries} attempts")
+                    raise
+                
+    def get_action(self, obs):
+
+        if self.batch_size == 0:
+            if obs["rgb_exo"].ndim == 3:
+                self.batch_size = 1
+            else:
+                self.batch_size = obs["rgb_exo"].shape[0]
+            self.action_buffer = [None] * self.batch_size
+            self.current_buffer_index = [0] * self.batch_size
+
+        if self.batch_size == 1:
+            obs = {key: np.expand_dims(value, axis=0) if isinstance(value, np.ndarray) else [value] for key, value in obs.items()}
+
+        
+        actions = []
+        for i in range(self.batch_size):
+
+            if self.action_buffer[i] is None or self.current_buffer_index[i] >= self.buffer_length:
+                model_input = {
+                    "observation/exterior_image_1_left": resize_with_pad(obs["rgb_exo"][i], 224, 224),
+                    "observation/wrist_image_left": resize_with_pad(obs["rgb_ego"][i], 224, 224),
+                    "observation/joint_position": np.array(obs["joint_positions"][i][:7]).reshape(7,),
+                    "observation/gripper_position": np.array(obs["grasp"][i]).reshape(1,),
+                    "prompt": f'pick up the {obs["object_name"][i]} and place it on a plate.',
+                }
+                self.action_buffer[i] = self.model.infer(model_input)["actions"]  
+                self.current_buffer_index[i] = 0    
+            current_action = self.action_buffer[i][self.current_buffer_index[i]]
+            self.current_buffer_index[i] += 1
+
+            gripper_pos = np.array([np.clip(current_action[7], 0.0, 1.0)])
+
+            if self.grasping_style == "binary":
+                gripper_pos = np.array([1.0]) if gripper_pos >= self.config["gripper_threshold"] else np.array([0.0])
+            elif self.grasping_style == "semi_binary":
+                gripper_pos = gripper_pos if gripper_pos <= self.config["gripper_threshold"] else np.array([1.0])
+            elif not self.grasping_style == "continuous":
+                raise ValueError(f"Invalid grasping style: {self.grasping_style}")
+                
+                
+            arm_output = current_action[:7].reshape(7)
+            actions.append(np.concatenate([arm_output, gripper_pos], axis=0).reshape(1, -1))
+
+        return np.array(actions).reshape(self.batch_size, 8)
+    
+    def reset(self, indicies=None):
+        if indicies is not None:
+            for i in indicies:
+                self.action_buffer[i] = None
+                self.current_buffer_index[i] = 0
+        else:
+            self.batch_size = 0
+            self.action_buffer = None
+            self.current_buffer_index = None
+
+
+def get_config():
+    config = {
+        "host": "localhost",
+        "port": 8000,
+        "gripper_threshold": 0.5,
+        "grasping_style": "binary",
+        "buffer_length": 8,
+    }
+    return config

baselines/rum/__init__.py

@@ -0,0 +1 @@
+from .policy import Policy

baselines/rum/loss_fns/__init__.py

@@ -0,0 +1,37 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import torch
+
+class BinarizeGripper(ABC):
+    def __init__(
+        self,
+        model: Optional[torch.nn.Module],
+        binarize_gripper: bool = False,
+        threshold: float = 0.5,
+        upper_value: float = 1.0,
+        lower_value: float = 0.0,
+        *args,
+        **kwargs,
+    ):
+        super().__init__()
+        assert (
+            not binarize_gripper or not model.relative_gripper
+        ), "Binarize gripper and relative gripper cannot be used together"
+
+        self.binarize_gripper = binarize_gripper
+        self.threshold = threshold
+        self.upper_value = upper_value
+        self.lower_value = lower_value
+
+    @abstractmethod
+    def binarize_gripper(self, actions):
+        # here the actions will be of shape (batch_size, seq_len, action_space)
+        # last element in action space is gripper values between 0 and 1
+
+        # binarize the gripper values
+        if self.binarize_gripper:
+            actions[:, :, -1] = torch.where(
+                actions[:, :, -1] > self.threshold, self.upper_value, self.lower_value
+            )
+        return actions

baselines/rum/loss_fns/abstract_loss_fn.py

@@ -0,0 +1,13 @@
+import abc
+
+import torch
+
+
+class AbstractLossFn(torch.nn.Module, abc.ABC):
+    def __init__(self, model, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        pass
+
+    @abc.abstractmethod
+    def forward(self, data, output, *args, **kwargs):
+        raise NotImplementedError

baselines/rum/loss_fns/diffusion_policy_loss_fn.py

@@ -0,0 +1,114 @@
+from typing import Optional, Sequence, Tuple
+
+import einops
+import torch
+import torch.nn as nn
+from itertools import chain
+
+from baselines.rum.loss_fns.abstract_loss_fn import AbstractLossFn
+from baselines.rum.models.policies.diffusion_policy import DiffusionPolicy
+
+
+class DiffusionPolicyLossFn(AbstractLossFn):
+    def __init__(
+        self,
+        tokenized_bet: bool,
+        action_dim: int,
+        obs_dim: int,
+        xyz_only: bool,
+        mask_last_min: int = 0,
+        mask_last_max: int = 0,
+        learned_mask: bool = True,
+        use_depth: bool = False,
+        model: Optional[torch.nn.Module] = None,
+        obs_window_size: int = 10,
+        action_sequence_length: int = 1,
+        data_act_scale: float = 1.0,
+        data_obs_scale: float = 1.0,
+        policy_type: str = "cnn",
+        device: str = "cuda",
+    ):
+        super().__init__(model)
+        assert mask_last_max >= mask_last_min
+        assert mask_last_min >= 0
+        obs_dim = model.feature_dim if not use_depth else model.feature_dim * 2
+        if use_depth:
+            self._depth_net = DepthNet(model.feature_dim)
+        self._use_depth = use_depth
+        self._true_action_dim = action_dim
+        action_dim = 3 if xyz_only else action_dim
+
+        self._diffusionpolicy = DiffusionPolicy(
+            obs_dim=obs_dim,
+            act_dim=action_dim,
+            obs_horizon=obs_window_size,
+            pred_horizon=(obs_window_size + action_sequence_length - 1),
+            action_horizon=action_sequence_length,
+            data_act_scale=data_act_scale,
+            data_obs_scale=data_obs_scale,
+            policy_type=policy_type,
+            device=device,
+        )
+
+        self._obs_mask_token = (
+            nn.Parameter(torch.ones(obs_dim), requires_grad=False)
+            if not learned_mask or mask_last_max == 0
+            else nn.Parameter(torch.randn(obs_dim), requires_grad=True)
+        )
+        self._obs_adapter = nn.Linear(obs_dim, obs_dim, bias=False)
+        self._adapt_obs = (
+            self._adapt_obs_linear if not tokenized_bet else self._adapt_obs_tokenized
+        )
+        self._mask_last = (mask_last_min, mask_last_max)
+        self._action_dim = action_dim
+
+        self.step = self._step if not tokenized_bet else self._step_tokenized
+
+        self._seen_action_stack = []
+        self._start_and_ends = None
+
+    def ema_step(self):
+        self._diffusionpolicy.ema_step()
+
+    def _adapt_obs_tokenized(self, obs):
+        return einops.rearrange(
+            self._obs_adapter(einops.rearrange(obs, "... c h w -> ... h w c")),
+            "... h w c -> ... c h w",
+        )
+
+    def _adapt_obs_linear(self, obs):
+        return self._obs_adapter(obs)
+
+    def _begin_epoch(self, *args, **kwargs):
+        return self._diffusionpolicy._begin_epoch(*args, **kwargs)
+
+    def forward(self, data, output, eval=False, *args, **kwargs):
+        *_, padding, actions = data
+        if self._use_depth:
+            *_, depths, padding, actions = data
+            output = torch.cat([output, self._depth_net(depths)], dim=-1)
+        adapted_obs = self._adapt_obs(output)
+        if actions is not None:
+            action_seq = actions[..., : self._action_dim].contiguous()
+        else:
+            action_seq = None
+        _, loss, loss_dict = self._diffusionpolicy(
+            adapted_obs, action_seq=action_seq, eval=eval
+        )
+        return loss, loss_dict
+
+    @torch.no_grad()
+    def _step(self, data, output, *args, **kwargs):
+        if self._use_depth:
+            _, depths, *_ = data
+            output = torch.cat([output, self._depth_net(depths)], dim=-1)
+        adapted_obs = self._adapt_obs(output)
+        a_hat, _, _ = self._diffusionpolicy(
+            adapted_obs,
+            action_seq=None,
+            eval=True,
+        )
+
+        return a_hat[0], {}
+
+        # return a_hat, {}

baselines/rum/loss_fns/rvq_loss_fn.py

@@ -0,0 +1,104 @@
+from typing import Optional, Sequence, Tuple
+
+import einops
+import torch
+import torch.nn as nn
+
+from baselines.rum.loss_fns.abstract_loss_fn import AbstractLossFn
+from baselines.rum.models.bet import (
+    GPT,
+    BehaviorTransformer,
+    TokenizedBehaviorTransformer,
+)
+from baselines.rum.models.bet.vqvae.vqvae import VqVae
+
+
+class RVQLossFn(AbstractLossFn):
+    def __init__(
+        self,
+        tokenized_bet: bool,
+        action_dim: int,
+        xyz_only: bool,
+        mask_last_min: int = 0,
+        mask_last_max: int = 0,
+        gpt_input_dim: int = 0,
+        learned_mask: bool = True,
+        model: Optional[torch.nn.Module] = None,
+        predict_with_offsets: bool = True,
+        sampling_temperature: float = 1.0,
+        action_sequence_length: int = 1,
+        vqvae_n_latent_dims: int = 512,
+        vqvae_n_embed: int = 16,
+        vqvae_groups: int = 2,
+        obs_cond: bool = False,
+    ):
+        super().__init__(model)
+        assert mask_last_max >= mask_last_min
+        assert mask_last_min >= 0
+        obs_dim = model.feature_dim
+        self._true_action_dim = action_dim
+        action_dim = 3 if xyz_only else action_dim
+        self.obs_cond = obs_cond
+        self._model = model
+        self._rvq = VqVae(
+            obs_dim=gpt_input_dim,
+            input_dim_h=action_sequence_length,
+            input_dim_w=action_dim,
+            n_latent_dims=vqvae_n_latent_dims,
+            vqvae_n_embed=vqvae_n_embed,
+            vqvae_groups=vqvae_groups,
+            eval=False,
+            device=model.device,
+            enc_loss_type="through_vqlayer",
+            obs_cond=obs_cond,
+        )
+
+        if self.obs_cond:
+            self._obs_mask_token = (
+                nn.Parameter(torch.ones(gpt_input_dim), requires_grad=False)
+                if not learned_mask or mask_last_max == 0
+                else nn.Parameter(torch.randn(gpt_input_dim), requires_grad=True)
+            )
+            self._obs_adapter = nn.Linear(obs_dim, gpt_input_dim, bias=False)
+            self._adapt_obs = (
+                self._adapt_obs_linear
+                if not tokenized_bet
+                else self._adapt_obs_tokenized
+            )
+        self._mask_last = (mask_last_min, mask_last_max)
+        self._action_dim = action_dim
+
+        self.step = self._step if not tokenized_bet else self._step_tokenized
+
+        self._seen_action_stack = []
+        self._start_and_ends = None
+
+    def _adapt_obs_tokenized(self, obs):
+        return einops.rearrange(
+            self._obs_adapter(einops.rearrange(obs, "... c h w -> ... h w c")),
+            "... h w c -> ... c h w",
+        )
+
+    def _adapt_obs_linear(self, obs):
+        return self._obs_adapter(obs)
+
+    def _begin_epoch(self, *args, **kwargs):
+        return self._rvq._begin_epoch(*args, **kwargs)
+
+    def forward(self, data, output, *args, **kwargs):
+        *_, padding, actions = data
+        action_seq = actions[..., : self._action_dim].contiguous()
+        if self.obs_cond:
+            adapted_obs = self._adapt_obs(output)
+            loss, loss_dict = self._rvq.vqvae_update(action_seq, adapted_obs)
+        else:
+            loss, loss_dict = self._rvq.vqvae_update(action_seq, None)
+        return loss, loss_dict
+
+    @torch.no_grad()
+    def _step(self, data, output, *args, **kwargs):
+        pass
+
+    @torch.no_grad()
+    def _step_tokenized(self, data, output, *args, **kwargs):
+        pass

baselines/rum/loss_fns/vqbet_loss_fn.py

@@ -0,0 +1,202 @@
+from typing import Optional
+
+import einops
+import torch
+import torch.nn as nn
+from einops.layers.torch import Rearrange
+from einops import rearrange
+
+from baselines.rum.loss_fns.abstract_loss_fn import AbstractLossFn
+from baselines.rum.models.bet import GPT
+from baselines.rum.models.bet.vqvae.vqvae import VqVae
+from baselines.rum.models.bet.vqbet import VQBehaviorTransformer
+
+
+class GoalAdapter(nn.Module):
+    def __init__(self, goal_dim, gpt_input_dim):
+        super().__init__()
+        self.linear = nn.Linear(goal_dim, gpt_input_dim, bias=False)
+
+    def forward(self, x):
+        # x: [b, t, g]
+        b, t, g = x.shape
+        x = rearrange(x, "b t g -> (b t) g")
+        x = self.linear(x)
+        x = rearrange(x, "(b t) g -> b t g", b=b, t=t)
+        return x
+
+
+class VQBeTLossFn(AbstractLossFn):
+    def __init__(
+        self,
+        tokenized_bet: bool,
+        action_dim: int,
+        xyz_only: bool,
+        vqvae_load_dir: str,
+        gpt_model: GPT,
+        goal_dim: int = 0,
+        mask_last_min: int = 0,
+        mask_last_max: int = 0,
+        learned_mask: bool = True,
+        use_depth: bool = False,
+        model: Optional[torch.nn.Module] = None,
+        action_sequence_length: int = 1,
+        vqvae_n_latent_dims: int = 512,
+        vqvae_n_embed: int = 16,
+        vqvae_groups: int = 2,
+        obs_cond: bool = False,
+        offset_loss_multiplier: float = 100.0,
+        secondary_code_multiplier: float = 0.5,
+        gamma: float = 2.0,
+        obs_window_size: int = 10,
+        sequentially_select: bool = False,
+        temperature: float = 1.0,
+        device: str = "cuda",
+    ):
+        super().__init__(model)
+        assert mask_last_max >= mask_last_min
+        assert mask_last_min >= 0
+        obs_dim = model.feature_dim if not use_depth else model.feature_dim * 2
+        gpt_input_dim = gpt_model.config.input_dim
+        if use_depth:
+            self._depth_net = DepthNet(model.feature_dim)
+        self._use_depth = use_depth
+        self.goal_dim = goal_dim
+
+        # TODO (mahi): currently, we are casting everything to a concat style goal
+        #  but we should be able to handle different types of goals like concat or stack
+        self._use_goals = goal_dim > 0
+        if self._use_goals:
+            gpt_input_dim //= 2
+            self._goal_adapter = GoalAdapter(goal_dim, gpt_input_dim)
+            goal_dim = gpt_input_dim
+        else:
+            self._goal_adapter = Rearrange("b g -> b 1 g")
+        self._true_action_dim = action_dim
+        action_dim = 3 if xyz_only else action_dim
+        self._rvq = VqVae(
+            obs_dim=gpt_input_dim,
+            input_dim_h=action_sequence_length,
+            input_dim_w=action_dim,
+            n_latent_dims=vqvae_n_latent_dims,
+            vqvae_n_embed=vqvae_n_embed,
+            vqvae_groups=vqvae_groups,
+            device=device,
+            eval=True,
+            enc_loss_type="through_vqlayer",
+            obs_cond=obs_cond,
+            load_dir=vqvae_load_dir,
+        )
+
+        for param in self._rvq.parameters():
+            param.requires_grad = False
+        self._vqbet = VQBehaviorTransformer(
+            obs_dim=gpt_input_dim,
+            act_dim=action_dim,
+            goal_dim=goal_dim,
+            gpt_model=gpt_model,
+            vqvae_model=self._rvq,
+            offset_loss_multiplier=offset_loss_multiplier,
+            secondary_code_multiplier=secondary_code_multiplier,
+            gamma=gamma,
+            obs_window_size=obs_window_size,
+            act_window_size=action_sequence_length,
+            sequentially_select=sequentially_select,
+            temperature=temperature,
+            device=device,
+        )
+
+        self._obs_mask_token = (
+            nn.Parameter(torch.ones(gpt_input_dim), requires_grad=False)
+            if not learned_mask or mask_last_max == 0
+            else nn.Parameter(torch.randn(gpt_input_dim), requires_grad=True)
+        )
+        self._obs_adapter = nn.Linear(obs_dim, gpt_input_dim, bias=False)
+        # self._adapt_obs = nn.Linear(obs_dim, gpt_input_dim, bias=False)
+        self._mask_last = (mask_last_min, mask_last_max)
+        self._action_dim = action_dim
+
+        self._adapt_obs = (
+            self._adapt_obs_linear if not tokenized_bet else self._adapt_obs_tokenized
+        )
+        self._mask_last = (mask_last_min, mask_last_max)
+        self._action_dim = action_dim
+
+        self.step = self._step if not tokenized_bet else self._step_tokenized
+
+        self._seen_action_stack = []
+        self._start_and_ends = None
+
+    def _adapt_obs_tokenized(self, obs):
+        return einops.rearrange(
+            self._obs_adapter(einops.rearrange(obs, "... c h w -> ... h w c")),
+            "... h w c -> ... c h w",
+        )
+
+    def _adapt_obs_linear(self, obs):
+        return self._obs_adapter(obs)
+
+    def _begin_epoch(self, *args, **kwargs):
+        return self._vqbet._begin_epoch(*args, **kwargs)
+
+    def forward(self, data, output, *args, **kwargs):
+        # TODO Mahi fix the order of depth and goals.
+        if self._use_goals:
+            _, goals, *_, padding, actions = data
+            if self._use_depth:
+                _, goals, *_, depths, padding, actions = data
+                output = torch.cat([output, self._depth_net(depths)], dim=-1)
+            goals = self._goal_adapter(goals)
+        else:
+            *_, padding, actions = data
+            goals = None
+            if self._use_depth:
+                *_, depths, padding, actions = data
+                output = torch.cat([output, self._depth_net(depths)], dim=-1)
+        adapted_obs = self._adapt_obs(output)
+        if "second_half" in kwargs:
+            second_half = kwargs["second_half"]
+        else:
+            second_half = False
+        action_seq = actions[..., : self._action_dim].contiguous()
+        _, loss, loss_dict = self._vqbet(
+            adapted_obs,
+            goal_seq=goals,
+            action_seq=action_seq,
+            second_half=second_half,
+        )
+        return loss, loss_dict
+
+    @torch.no_grad()
+    def _step(self, data, output, return_all=False, *args, **kwargs):
+        if self._use_depth:
+            *_, depths, padding, actions = data
+            output = torch.cat([output, self._depth_net(depths)], dim=-1)
+        goals = data[1] if self._use_goals else None
+        adapted_obs = self._adapt_obs(output)
+        goals = self._goal_adapter(goals) if self._use_goals else None
+        a_hat, _, _ = self._vqbet(
+            adapted_obs,
+            goal_seq=goals,
+            action_seq=None,
+        )
+        if a_hat.shape[2] != self._true_action_dim:
+            # append n zeros and a 1 to the end of y_hat
+            a_hat = torch.cat(
+                [
+                    a_hat,
+                    torch.zeros(
+                        a_hat.shape[0],
+                        a_hat.shape[1],
+                        self._true_action_dim - a_hat.shape[2],
+                    ).to(a_hat.device),
+                ],
+                dim=2,
+            )
+            a_hat[:, :, -1] = 1.0
+        if return_all:
+            return a_hat.reshape(
+                adapted_obs.shape[0], adapted_obs.shape[1], self._true_action_dim
+            )[:, -1, :], {}
+        # Finally, return the final action prediction only.
+        return a_hat[-1, -1, :], {}

baselines/rum/models/__init__.py

@@ -0,0 +1 @@
+from . import *

baselines/rum/models/bet/__init__.py

@@ -0,0 +1,3 @@
+from baselines.rum.models.bet.bet import BehaviorTransformer
+from baselines.rum.models.bet.gpt import GPT
+from baselines.rum.models.bet.tokenized_bet import TokenizedBehaviorTransformer