egogym 0.1.0__py3-none-any.whl

baselines/rum/utils/__init__.py

@@ -0,0 +1,212 @@
+import random
+from collections import OrderedDict
+from datetime import timedelta
+
+import accelerate
+import numpy as np
+import torch
+import json
+
+from . import *
+
+# import accelerate
+
+
+def set_seed_everywhere(seed):
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+
+def get_demos(cfg):
+    train_files_list = cfg.dataset.train.config.trajectory_roots.trajectory_root_path
+    val_files_list = cfg.dataset.test.config.trajectory_roots.trajectory_root_path
+
+    train_demos = json.load(open(train_files_list))
+    val_demos = json.load(open(val_files_list))
+
+    return {"train": train_demos, "val": val_demos}
+
+
+def setup_accelerate(experiment, num_epochs, task_title, save_every):
+    ddp_kwargs = accelerate.DistributedDataParallelKwargs(find_unused_parameters=True)
+
+    accelerator = accelerate.Accelerator(
+        log_with=["wandb"],
+        kwargs_handlers=[
+            accelerate.InitProcessGroupKwargs(timeout=timedelta(hours=1.5)),
+            ddp_kwargs,
+        ],
+    )
+    accelerator.init_trackers(
+        "experiment",
+        config={
+            "num_epochs": 1000,
+            "task_title": task_title,
+            "save_every": 20,
+        },
+    )
+
+    return accelerator, ddp_kwargs
+
+
+class AverageMeter:
+    def __init__(self, value=None):
+        super().__init__()
+        if value is None:
+            self._sum = 0
+            self._count = 0
+        else:
+            self._sum = value
+            self._count = 1
+
+    def update(self, value, n=1):
+        self._sum += value * n
+        self._count += n
+        self._avg = self._sum / self._count
+
+    @property
+    def avg(self):
+        return self._avg
+
+
+class AveragingDict(OrderedDict):
+    def __init__(self, name="train", *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.name = name
+
+    def update(self, other):
+        for k, v in other.items():
+            if k not in self:
+                self[k] = AverageMeter(v)
+            self[k].update(v)
+
+    def __str__(self):
+        return ", ".join(f"{self.name}/{k}: {v.avg:.2E}" for k, v in self.items())
+
+    @property
+    def full_summary(self):
+        return {f"{self.name}/{k}": v.avg for k, v in self.items()}
+
+    @property
+    def summary(self):
+        return {f"{self.name}/{k}": f"{v.avg:.2E}" for k, v in self.items()}
+
+
+class Callbacks:
+    """
+    This class is used to implement callback functions inside training framework that
+    passes, training params to callback functions inside model and loss objects.
+
+    This is useful for custom logging inside models and losses like adding image
+    visualisation, plotting hyperparameters like learning rate, etc.
+    """
+
+    def __init__(self):
+        pass
+
+    def set_workspace(self, workspace, accelerator):
+        self.workspace = workspace
+        self.accelerator = accelerator
+        self._model_has_begin_epoch = hasattr(
+            self.accelerator.unwrap_model(self.workspace.model), "_begin_epoch"
+        )
+        self._model_has_begin_batch = hasattr(
+            self.accelerator.unwrap_model(self.workspace.model), "_begin_batch"
+        )
+
+        self._loss_fn_has_begin_epoch = hasattr(
+            self.accelerator.unwrap_model(self.workspace.loss_fn), "_begin_epoch"
+        )
+        self._loss_fn_has_begin_batch = hasattr(
+            self.accelerator.unwrap_model(self.workspace.loss_fn), "_begin_batch"
+        )
+
+    def begin_epoch(self):
+        self.batch_step = 0
+
+        if self._model_has_begin_epoch:
+            model = self.accelerator.unwrap_model(self.workspace.model)
+            is_main_process = self.accelerator.is_local_main_process
+            begin_epoch_log = model._begin_epoch(
+                epoch=self.workspace._epoch,
+                epochs=self.workspace._cfg.num_epochs,
+                train_dataloader=self.workspace._train_dataloaders,
+                test_dataloader=self.workspace._test_dataloaders,
+                optimizer=self.workspace.optimizer,
+                scheduler=self.workspace.scheduler,
+                is_main_process=is_main_process,
+            )
+            if begin_epoch_log is not None:
+                self.accelerator.log(begin_epoch_log)
+
+        if self._loss_fn_has_begin_epoch:
+            loss_fn = self.accelerator.unwrap_model(self.workspace.loss_fn)
+            is_main_process = self.accelerator.is_local_main_process
+            begin_epoch_log = loss_fn._begin_epoch(
+                epoch=self.workspace._epoch,
+                epochs=self.workspace._cfg.num_epochs,
+                train_dataloader=self.workspace._train_dataloaders,
+                test_dataloader=self.workspace._test_dataloaders,
+                optimizer=self.workspace.optimizer,
+                scheduler=self.workspace.scheduler,
+                is_main_process=is_main_process,
+            )
+            if begin_epoch_log is not None:
+                self.accelerator.log(begin_epoch_log)
+
+    def begin_batch(self):
+        if self._model_has_begin_batch:
+            model = self.accelerator.unwrap_model(self.workspace.model)
+            is_main_process = self.accelerator.is_local_main_process
+            begin_batch_log = model._begin_batch(
+                epoch=self.workspace._epoch,
+                batch_step=self.batch_step,
+                epochs=self.workspace._cfg.num_epochs,
+                train_dataloader=self.workspace._train_dataloaders,
+                test_dataloader=self.workspace._test_dataloaders,
+                optimizer=self.workspace.optimizer,
+                scheduler=self.workspace.scheduler,
+                is_main_process=is_main_process,
+            )
+            if begin_batch_log is not None:
+                self.accelerator.log(begin_batch_log)
+
+        if self._loss_fn_has_begin_batch:
+            loss_fn = self.accelerator.unwrap_model(self.workspace.loss_fn)
+            is_main_process = self.accelerator.is_local_main_process
+            begin_batch_log = loss_fn._begin_batch(
+                epoch=self.workspace._epoch,
+                batch_step=self.batch_step,
+                epochs=self.workspace._cfg.num_epochs,
+                train_dataloader=self.workspace._train_dataloaders,
+                test_dataloader=self.workspace._test_dataloaders,
+                optimizer=self.workspace.optimizer,
+                scheduler=self.workspace.scheduler,
+                is_main_process=is_main_process,
+            )
+            if begin_batch_log is not None:
+                self.accelerator.log(begin_batch_log)
+
+        self.batch_step += 1
+
+
+def log_wandb_image(image, label):
+    """
+    image: torch.Tensor or np.ndarray
+    label: str, label for the image to categorize it in wandb
+
+    returns: dict, {label: wandb.Image(image)}
+    """
+    import wandb
+
+    if type(image) == torch.Tensor:
+        image = image.cpu().numpy()
+    while image.shape[0] == 1:
+        image = image.squeeze(0)
+    if image.shape[0] == 3:
+        image = image.transpose(1, 2, 0)
+    image = wandb.Image(image)
+    return {label: image}

baselines/rum/utils/action_transforms.py

@@ -0,0 +1,22 @@
+import numpy as np
+
+# permutation transformation matrix to go from record3d axis to personal camera axis
+P = np.array([[-1, 0, 0, 0], [0, 0, -1, 0], [0, -1, 0, 0], [0, 0, 0, 1]])
+# end effector transformation matrix to go from personal camera axis to end effector axis. Corrects the 15 degree offset along the x axis
+EFT = np.asarray([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
+
+
+def apply_permutation_transform(matrix):
+    return P @ matrix @ P.T
+
+
+def invert_permutation_transform(matrix):
+    return P.T @ matrix @ P
+
+
+def apply_end_effector_transform(matrix):
+    return EFT @ matrix
+
+
+def invert_end_effector_transform(matrix):
+    return EFT.T @ matrix

baselines/rum/utils/decord_transforms.py

@@ -0,0 +1,135 @@
+import logging
+from typing import Iterable, Optional, Tuple, Union
+
+import torch
+from timm.data.random_erasing import RandomErasing
+from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+
+DEFAULT_CROP_PCT = 0.875
+DEFAULT_CROP_MODE = "center"
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+IMAGENET_INCEPTION_MEAN = (0.5, 0.5, 0.5)
+IMAGENET_INCEPTION_STD = (0.5, 0.5, 0.5)
+IMAGENET_DPN_MEAN = (124 / 255, 117 / 255, 104 / 255)
+IMAGENET_DPN_STD = tuple([1 / (0.0167 * 255)] * 3)
+OPENAI_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)
+
+
+def create_transform(
+    input_size,
+    is_training: bool = False,
+    scale: Optional[Tuple[float, float]] = None,
+    ratio: Optional[Tuple[float, float]] = None,
+    hflip: float = 0.5,
+    vflip: float = 0.0,
+    grayscale: float = 0.0,
+    gaussblr: float = 0.0,
+    gaussblr_kernel: int = 3,
+    gaussblr_sigma: tuple = (1.0, 2.0),
+    color_jitter: Optional[Union[Iterable[float], float]] = 0.4,
+    interpolation: Union[str, InterpolationMode] = "bilinear",
+    mean: Union[Iterable, torch.Tensor] = IMAGENET_DEFAULT_MEAN,
+    std: Union[Iterable, torch.Tensor] = IMAGENET_DEFAULT_STD,
+    re_prob=0.0,
+    re_mode="const",
+    re_count=1,
+    re_num_splits=0,
+    crop_pct=None,
+    crop_mode=None,
+    *args,
+    **kwargs,
+):
+    if isinstance(input_size, (tuple, list)):
+        img_size = input_size[-2:]
+    else:
+        img_size = input_size
+    transform_list = []
+
+    interpolation_mode = InterpolationMode.BILINEAR
+    try:
+        interpolation_mode = InterpolationMode[interpolation.upper()]
+    except AttributeError:
+        logging.warning(
+            f"Interpolation mode {interpolation} is not recognized, "
+            f"using bilinear instead."
+        )
+    if is_training:
+        scale = tuple(scale or (0.08, 1.0))  # default imagenet scale range
+        ratio = tuple(ratio or (3.0 / 4.0, 4.0 / 3.0))  # default imagenet ratio range
+
+        transform_list.append(
+            transforms.RandomResizedCrop(
+                img_size,
+                scale,
+                ratio,
+                antialias=False,
+                interpolation=interpolation_mode,
+            )
+        )
+        if hflip > 0.0:
+            transform_list.append(transforms.RandomHorizontalFlip(hflip))
+        if vflip > 0.0:
+            transform_list.append(transforms.RandomVerticalFlip(vflip))
+
+        if color_jitter is not None:
+            if isinstance(color_jitter, (list, tuple)):
+                assert len(color_jitter) in (
+                    3,
+                    4,
+                ), "expected either 3 or 4 values for color jitter"
+            else:
+                color_jitter = (float(color_jitter),) * 3
+            transform_list.append(transforms.ColorJitter(*color_jitter))
+
+        if grayscale > 0.0:
+            transform_list.append(transforms.RandomGrayscale(p=grayscale))
+        if gaussblr > 0.0:
+            transform_list.append(
+                transforms.RandomApply(
+                    [
+                        transforms.GaussianBlur(
+                            (gaussblr_kernel, gaussblr_kernel), gaussblr_sigma
+                        )
+                    ],
+                    p=gaussblr,
+                )
+            )
+        transform_list.append(transforms.Normalize(mean, std))
+
+        if re_prob > 0.0:
+            transform_list.append(
+                RandomErasing(
+                    re_prob,
+                    mode=re_mode,
+                    max_count=re_count,
+                    num_splits=re_num_splits,
+                    device="cuda",
+                )
+            )
+
+    else:
+        if crop_pct is None:
+            crop_pct = DEFAULT_CROP_PCT
+        if crop_mode is None:
+            crop_mode = DEFAULT_CROP_MODE
+        rescaled_size = (
+            int(img_size / crop_pct)
+            if isinstance(img_size, (int, float))
+            else ([int(size / crop_pct) for size in img_size])
+        )
+        # TODO (haritheja): we are no longer cropping, but figure this out
+        # transform_list.append(
+        #     transforms.Resize(rescaled_size, interpolation_mode, antialias=False)
+        # )
+        # if crop_mode == "center":
+        #     transform_list.append(transforms.CenterCrop(img_size))
+        # elif crop_mode == "random":
+        #     transform_list.append(transforms.RandomCrop(img_size))
+        # else:
+        #     raise ValueError(f"crop_mode '{crop_mode}' not recognized")
+        transform_list.append(transforms.Normalize(mean, std))
+
+    return transforms.Compose(transform_list)

baselines/rum/utils/rpc.py

@@ -0,0 +1,249 @@
+import zmq
+import time
+import pickle
+import threading
+import traceback
+
+
+class RPCServer:
+    def __init__(self, obj, port: int = 5000, threaded=False):
+        """
+        obj: object with methods to expose
+        port: port to listen on
+        """
+        self.obj = obj
+        self.context = zmq.Context()
+        self.socket: zmq.socket.Socket = self.context.socket(zmq.REP)
+        self.socket.bind(f"tcp://*:{port}")
+        self.threaded = threaded
+        if threaded:
+            self.thread = threading.Thread(target=self.run)
+            self.stop_event = threading.Event()
+        else:
+            self.stop_event = False
+
+    def _send_exception(self, e):
+        """
+        Serialize an exception and send it over the socket.
+        Only the exception type, message, and traceback are sent.
+        """
+        exception = {
+            "type": "exception",
+            "content": {
+                "exception": str(type(e)),
+                "message": str(e),
+                "traceback": traceback.format_exc(),
+            },
+        }
+        self.socket.send(pickle.dumps(exception))
+
+    def _send_result(self, result):
+        """
+        Serialize a result and send it over the socket.
+        """
+        result = {"type": "result", "content": result}
+        self.socket.send(pickle.dumps(result))
+
+    def run(self):
+        """
+        Run the server.
+        """
+        if self.threaded:
+            while not self.stop_event.is_set():
+                message = self.socket.recv()
+                message = pickle.loads(message)
+                self._handle_message(message)
+        else:
+            while not self.stop_event:
+                try:
+                    message = self.socket.recv(flags=zmq.NOBLOCK)
+                    message = pickle.loads(message)
+                except zmq.Again:
+                    time.sleep(0.001)
+                    continue
+                self._handle_message(message)
+
+    def _is_callable(self, attr):
+        return hasattr(self.obj, attr) and callable(getattr(self.obj, attr))
+
+    def _handle_message(self, message):
+        """
+        Handles a dictionary of {
+            "req": str,  # request type
+            "attr": str,
+            "args": list,
+            "kwargs": dict,
+        }
+        from the socket.
+        If req == "is_callable", return whether the attribute is callable.
+        If req == "get", return the attribute.
+            If the attribute is not found, return an error message.
+            If the attribute is callable, call with args and kwargs.
+                If there are any errors in the callable, return the pickled error
+                If the callable is found and there are no errors, return the pickled result.
+            If the attribute is not callable, return the attribute.
+        If req == "set", set the attribute to the value.
+        If req == "dir", return a list of attributes.
+        If req == "stop", stop the server.
+        """
+        if message["req"] == "is_callable":
+            result = self._is_callable(message["attr"])
+            self._send_result(result)
+        elif message["req"] == "get":
+            try:
+                attribute = getattr(self.obj, message["attr"])
+                args = message["args"]
+                kwargs = message["kwargs"]
+                if not callable(attribute):
+                    self._send_result(attribute)
+                else:
+                    result = attribute(*args, **kwargs)
+                    self._send_result(result)
+            except Exception as e:
+                self._send_exception(e)
+        elif message["req"] == "set":
+            try:
+                setattr(self.obj, message["attr"], message["value"])
+                self._send_result(None)
+            except Exception as e:
+                self._send_exception(e)
+        elif message["req"] == "dir":
+            result = dir(self.obj)
+            self._send_result(result)
+        elif message["req"] == "stop":
+            self.stop()
+
+    def close(self):
+        self.socket.close()
+        self.context.term()
+
+    def start(self):
+        if self.threaded:
+            self.stop_event.clear()
+            self.thread.start()
+        else:
+            self.run()
+
+    def stop(self):
+        if self.threaded:
+            self.stop_event.set()
+            self.thread.join()
+        else:
+            self.stop_event = True
+        self.close()
+
+
+class RPCException(Exception):
+    def __init__(self, exception_type: str, message: str, traceback: str):
+        self.exception_type = exception_type
+        self.message = message
+        self.traceback = traceback
+
+    def __str__(self):
+        return f"{self.exception_type}: {self.message}\n{self.traceback}"
+
+
+class RPCClient:
+    def __init__(self, host: str, port: int = 5000):
+        """
+        host: host to connect to
+        port: port to connect to
+        """
+        self.__dict__["context"] = zmq.Context()
+        self.__dict__["socket"] = self.context.socket(zmq.REQ)
+        self.socket.connect(f"tcp://{host}:{port}")
+        self.__dict__["_is_callable_cache"] = {}
+
+    def __setattr__(self, attr: str, value):
+        """
+        Set the attribute of the same name.
+        Attribute must not be callable on remote.
+        """
+        if self._is_callable(attr):
+            raise AttributeError(f"Overwriting a callable attribute: {attr}")
+        self._send_set(attr, value)
+
+    def _send_get(self, attr: str, args: list, kwargs: dict):
+        """
+        Send a get request over the socket.
+        """
+        req = {"req": "get", "attr": attr, "args": args, "kwargs": kwargs}
+        self.socket.send(pickle.dumps(req))
+        return self._recv_result()
+
+    def _send_set(self, attr: str, value):
+        """
+        Send a set request over the socket.
+        """
+        req = {"req": "set", "attr": attr, "value": value}
+        self.socket.send(pickle.dumps(req))
+        return self._recv_result()
+
+    def _recv_result(self):
+        """
+        Receive a dictionary of {
+            "type": str,
+            "content": object,
+        }
+        if type == "exception", content is a dictionary of {
+            "exception": str,
+            "message": str,
+            "traceback": str,
+        }; re-raise the exception on the client side
+        if type == "result", content is the result
+        """
+        result = self.socket.recv()
+        result = pickle.loads(result)
+        if result["type"] == "exception":
+            raise RPCException(
+                result["content"]["exception"],
+                result["content"]["message"],
+                result["content"]["traceback"],
+            )
+        return result["content"]
+
+    def _is_callable(self, attr: str) -> bool:
+        """
+        Send a request to check if the attribute is callable.
+        Returns False if the attribute is not found.
+        """
+        if attr not in self._is_callable_cache:
+            req = {"req": "is_callable", "attr": attr}
+            self.socket.send(pickle.dumps(req))
+            result = self._recv_result()
+            self._is_callable_cache[attr] = result
+        return self._is_callable_cache[attr]
+
+    def __getattr__(self, attr: str):
+        """
+        Return the attribute of the same name.
+        If the attribute is a callable, return a function that sends the call over the socket.
+        Else, return the attribute value.
+        """
+        if self._is_callable(attr):
+            return lambda *args, **kwargs: self._send_get(attr, args, kwargs)
+        else:
+            return self._send_get(attr, [], {})
+
+    def __dir__(self):
+        """
+        Return a list of attributes.
+        """
+        req = {"req": "dir"}
+        self.socket.send(pickle.dumps(req))
+        result = self._recv_result()
+        return result + ["stop_server"]
+
+    def stop_server(self) -> bool:
+        """
+        Send a stop request to the server.
+        If the server is stopped, close the socket and terminate the context.
+        Returns a bool for success.
+        """
+        req = {"req": "stop"}
+        self.socket.send(pickle.dumps(req))
+        stopped = self._recv_result()
+        if stopped:
+            self.socket.close()
+            self.context.term()
+        return stopped

baselines/rum/utils/schedulers.py

@@ -0,0 +1,71 @@
+import math
+import warnings
+
+from torch.optim.lr_scheduler import LRScheduler
+
+
+class CosineAnnealWithWarmupLR(LRScheduler):
+    """Sets the learning rate of each parameter group to the initial lr
+    times a given function. When last_epoch=-1, sets initial lr as lr.
+
+    Adapted from karpathy/nanoGPT
+    https://github.com/karpathy/nanoGPT/blob/eba36e84649f3c6d840a93092cb779a260544d08/train.py#L228
+
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        warmup_epochs (int): The number of epochs for warmup
+        lr_decay_epoch (int): The index of lr_decay epoch, at which we reach min_LR. Default: -1.
+        verbose (bool): If ``True``, prints a message to stdout for
+            each update. Default: ``False``.
+        min_lr_multiplier: The lowest lr multiplier the schedule will ever get to.
+
+    Example:
+        >>> warmup_epochs = 10
+        >>> lr_decay_epochs = 600
+        >>> scheduler = CosineAnnealWithWarmupLR(optimizer, warmup_epochs, lr_decay_epochs)
+        >>> for epoch in range(100):
+        >>>     train(...)
+        >>>     validate(...)
+        >>>     scheduler.step()
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        warmup_epochs,
+        lr_decay_epochs,
+        last_epoch=-1,
+        min_lr_multiplier=0.1,
+        verbose=False,
+    ):
+        self.optimizer = optimizer
+        self.warmup_epochs = warmup_epochs
+        self.last_epoch = last_epoch
+        self.lr_decay_epoch = lr_decay_epochs
+        self.min_lr_multiplier = min_lr_multiplier
+
+        super().__init__(optimizer, last_epoch, verbose)
+
+    def get_lr(self):
+        if not self._get_lr_called_within_step:
+            warnings.warn(
+                "To get the last learning rate computed by the scheduler, "
+                "please use `get_last_lr()`."
+            )
+
+        return [
+            base_lr * self._calculate_lr_formula(self.last_epoch)
+            for base_lr in self.base_lrs
+        ]
+
+    def _calculate_lr_formula(self, last_epoch):
+        if last_epoch < self.warmup_epochs:
+            return (last_epoch + 1) / self.warmup_epochs
+        elif last_epoch > self.lr_decay_epoch:
+            return self.min_lr_multiplier
+        decay_ratio = (last_epoch - self.warmup_epochs) / (
+            self.lr_decay_epoch - self.warmup_epochs
+        )
+        assert 0 <= decay_ratio <= 1
+        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # coeff ranges 0..1
+        return self.min_lr_multiplier + coeff * (1 - self.min_lr_multiplier)