robot-utils 0.1.1__tar.gz

robot_utils-0.1.1/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: robot-utils
+Version: 0.1.1
+Summary: Frequently used utility functions for robot learning research
+Author-email: Yihuai Gao <davidgao1013@gmail.com>
+Project-URL: Homepage, https://github.com/yihuai-gao/robot-utils
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent

robot_utils-0.1.1/pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "robot-utils"
+version = "0.1.1"
+description = "Frequently used utility functions for robot learning research"
+authors = [{ name = "Yihuai Gao", email = "davidgao1013@gmail.com" }]
+dependencies = []
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = [
+    "robot_utils*",
+]
+
+[project.urls]
+Homepage = "https://github.com/yihuai-gao/robot-utils"

robot_utils-0.1.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

robot_utils-0.1.1/src/robot_utils/config_utils.py

@@ -0,0 +1,43 @@
+import hydra
+from omegaconf import DictConfig, ListConfig, OmegaConf
+from typing import Any, cast
+
+def register_resolvers():
+    OmegaConf.register_new_resolver("eval", eval)
+    OmegaConf.register_new_resolver("concat", lambda *args: args[0] + args[1])
+    OmegaConf.register_new_resolver("cond", lambda cond, true_val, false_val: true_val if cond else false_val)
+    OmegaConf.register_new_resolver("range", lambda start, end, step=1: ListConfig(list(range(start, end, step))))
+    OmegaConf.register_new_resolver("in", lambda target, *items: any(item in target for item in items))
+
+def disable_hydra_target(cfg: Any) -> Any:
+    if isinstance(cfg, DictConfig):
+        return_dict = DictConfig({})
+        for key, value in cfg.items():
+            if key == "_target_":
+                return_dict["__target__"] = value
+            else:
+                return_dict[key] = disable_hydra_target(value)
+        return return_dict
+    elif isinstance(cfg, ListConfig):
+        return_list = ListConfig([])
+        for item in cfg:
+            return_list.append(disable_hydra_target(item))
+        return return_list
+    return cfg
+
+
+def enable_hydra_target(cfg: Any) -> Any:
+    if isinstance(cfg, DictConfig) or isinstance(cfg, dict):
+        return_dict = DictConfig({})
+        for key, value in cfg.items():
+            if key == "__target__":
+                return_dict["_target_"] = value
+            else:
+                return_dict[key] = enable_hydra_target(value)
+        return return_dict
+    elif isinstance(cfg, ListConfig) or isinstance(cfg, list):
+        return_list = ListConfig([])
+        for item in cfg:
+            return_list.append(enable_hydra_target(item))
+        return return_list
+    return cfg

robot_utils-0.1.1/src/robot_utils/data_utils.py

@@ -0,0 +1,172 @@
+import numpy as np
+import numpy.typing as npt
+import cv2
+from typing import Any, Callable, Union
+from cv2.typing import MatLike
+
+
+def resize_with_padding(
+    img: npt.NDArray[Any], new_shape_hw: tuple[int, ...]
+) -> npt.NDArray[Any]:
+    """
+    img: (..., C, H, W)
+
+    new_shape_hw: (new_H, new_W)
+
+    return: (..., C, new_H, new_W)
+    """
+
+    assert (
+        len(new_shape_hw) == 2
+    ), f"new_shape_hw must be a tuple of length 2, but got {new_shape_hw}"
+
+    batch_shape = img.shape[:-3]
+    C, H, W = img.shape[-3:]
+
+    original_aspect_ratio = W / H
+    new_H = new_shape_hw[0]
+    new_W = new_shape_hw[1]
+    new_aspect_ratio = new_W / new_H
+    if original_aspect_ratio > new_aspect_ratio:
+        # Pad upwards and downwards
+        new_H_without_padding = int(new_W / original_aspect_ratio)
+        new_W_without_padding = new_W
+        padding_top = (new_H - new_H_without_padding) // 2
+        padding_bottom = new_H - new_H_without_padding - padding_top
+        padding_sequence = [(0, 0) for _ in range(len(batch_shape))] + [
+            (0, 0),
+            (padding_top, padding_bottom),
+            (0, 0),
+        ]
+
+    else:
+        # Pad left and right
+        new_W_without_padding = int(new_H * original_aspect_ratio)
+        new_H_without_padding = new_H
+        padding_left = (new_W - new_W_without_padding) // 2
+        padding_right = new_W - new_W_without_padding - padding_left
+        padding_sequence = [(0, 0) for _ in range(len(batch_shape))] + [
+            (0, 0),
+            (0, 0),
+            (padding_left, padding_right),
+        ]
+
+    img = img.reshape(-1, 3, H, W)
+    resized_img = np.zeros(
+        (img.shape[0], 3, new_H_without_padding, new_W_without_padding), dtype=img.dtype
+    )
+    for i in range(img.shape[0]):
+        img_hwc = img[i].transpose(1, 2, 0)
+        img_hwc = cv2.resize(img_hwc, (new_W_without_padding, new_H_without_padding))
+        resized_img[i] = img_hwc.transpose(2, 0, 1)
+    resized_img = resized_img.reshape(
+        *batch_shape, C, new_H_without_padding, new_W_without_padding
+    )
+    resized_img = np.pad(
+        resized_img, padding_sequence, mode="constant", constant_values=0
+    )
+
+    return resized_img
+
+
+def resize_with_cropping(
+    source_frame_hwc: npt.NDArray[Any],
+    display_wh: tuple[int, int],
+) -> npt.NDArray[Any]:
+    """
+    source_frame: (..., H, W, C)
+    Crops and resizes a source frame to fit the display resolution while preserving aspect ratio.
+    This logic is adapted directly from the user's robust BaseCamera implementation.
+    """
+    source_height, source_width = source_frame_hwc.shape[-3:-1]
+    display_width, display_height = display_wh
+
+    source_wh_ratio = source_width / source_height
+    display_wh_ratio = display_width / display_height
+
+    if source_wh_ratio > display_wh_ratio:
+        # source is "wider" than display, crop the width
+        new_width = int(source_height * display_wh_ratio)
+        margin = (source_width - new_width) // 2
+        cropped_frame = source_frame_hwc[..., :, margin : margin + new_width, :]
+    else:
+        # source is "taller" than or same as display, crop the height
+        new_height = int(source_width / display_wh_ratio)
+        margin = (source_height - new_height) // 2
+        cropped_frame = source_frame_hwc[..., margin : margin + new_height, :, :]
+
+    if len(source_frame_hwc.shape) == 4:
+        resized_images = np.zeros(
+            (
+                source_frame_hwc.shape[0],
+                display_height,
+                display_width,
+                source_frame_hwc.shape[-1],
+            ),
+            dtype=source_frame_hwc.dtype,
+        )
+        for i in range(source_frame_hwc.shape[0]):
+            resized_images[i] = cv2.resize(cropped_frame[i], display_wh)
+    else:
+        resized_images = cv2.resize(cropped_frame, display_wh)
+    resized_images = np.array(resized_images)
+    return resized_images
+
+
+
+def aggregate_dict(
+    dictionaries: list[dict[str, Any]] | dict[Any, dict[str, Any]],
+    convert_to_numpy: bool,
+    key_name: str = "",
+) -> dict[str, Any]:
+    """
+    Aggregate a list of dictionaries or a dictionary of dictionaries into a single dictionary.
+    """
+    aggregated_dict: dict[str, Any] = {}
+    if isinstance(dictionaries, list):
+        for single_dict in dictionaries:
+            for key, value in single_dict.items():
+                if key not in aggregated_dict:
+                    aggregated_dict[key] = []
+                aggregated_dict[key].append(value)
+
+    elif isinstance(dictionaries, dict):
+        assert key_name != "", "Key name is required for dictionary aggregation"
+        aggregated_dict[key_name] = []
+        for key, value in dictionaries.items():
+            assert (
+                key_name not in value
+            ), f"Key {key_name} is not allowed in the dictionary. Please rename the key for aggregation"
+            aggregated_dict[key_name].append(key)
+            for key, value in value.items():
+                if key not in aggregated_dict:
+                    aggregated_dict[key] = []
+                aggregated_dict[key].append(value)
+
+    for key, value in aggregated_dict.items():
+        if key == key_name:
+            continue
+        if isinstance(value, list):
+            if isinstance(value[0], dict):
+                # Value is a list of dictionaries. Will be flattened in the next step
+                value = aggregate_dict(value, convert_to_numpy)
+            elif convert_to_numpy:
+                if not isinstance(value[0], str):
+                    try:
+                        aggregated_dict[key] = np.array(value)
+                    except Exception as e:
+                        print(f"Error aggregating {key}: {e}")
+                        for v in value:
+                            print(f"{v.shape}")
+                        raise e
+    return aggregated_dict
+
+
+def dict_apply(x: dict[str, Any], func: Callable[[Any], Any]) -> dict[str, Any]:
+    result = dict()
+    for key, value in x.items():
+        if isinstance(value, dict):
+            result[key] = dict_apply(value, func)
+        else:
+            result[key] = func(value)
+    return result

robot_utils-0.1.1/src/robot_utils/logging_utils.py

@@ -0,0 +1,10 @@
+import sys
+import traceback
+
+
+def echo_exception():
+    exc_type, exc_value, exc_traceback = sys.exc_info()
+    # Extract unformatted traceback
+    tb_lines = traceback.format_exception(exc_type, exc_value, exc_traceback)
+    # Print line of code where the exception occurred
+    return "".join(tb_lines)

robot_utils-0.1.1/src/robot_utils/pose_utils.py

@@ -0,0 +1,286 @@
+from typing import Any
+import numpy as np
+import numpy.typing as npt
+import scipy.spatial.transform as st
+from scipy.spatial.transform import Rotation
+
+
+def qmult(q1: npt.NDArray[Any], q2: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    q = np.array(
+        [
+            q1[0] * q2[0] - q1[1] * q2[1] - q1[2] * q2[2] - q1[3] * q2[3],
+            q1[0] * q2[1] + q1[1] * q2[0] + q1[2] * q2[3] - q1[3] * q2[2],
+            q1[0] * q2[2] - q1[1] * q2[3] + q1[2] * q2[0] + q1[3] * q2[1],
+            q1[0] * q2[3] + q1[1] * q2[2] - q1[2] * q2[1] + q1[3] * q2[0],
+        ]
+    )
+
+    return q
+
+
+def qconjugate(q: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    return np.array([q[0], -q[1], -q[2], -q[3]])
+
+
+def to_xyzw(wxyz: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+    if wxyz.ndim == 1:
+        return np.concatenate([wxyz[1:], wxyz[0:1]])
+    elif wxyz.ndim == 2:
+        return np.concatenate([wxyz[:, 1:], wxyz[:, 0:1]], axis=1)
+    else:
+        raise ValueError("wxyz must be a 1D or 2D array")
+
+
+def to_wxyz(xyzw: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+    if xyzw.ndim == 1:
+        return np.concatenate([xyzw[3:], xyzw[:3]])
+    elif xyzw.ndim == 2:
+        return np.concatenate([xyzw[:, 3:], xyzw[:, :3]], axis=1)
+    else:
+        raise ValueError("xyzw must be a 1D or 2D array")
+
+
+def interpolate_xyz_wxyz(
+    pose_left: npt.NDArray[np.float64],
+    pose_right: npt.NDArray[np.float64],
+    timestamp_left: float,
+    timestamp_right: float,
+    timestamp: float,
+) -> npt.NDArray[np.float64]:
+    assert (
+        pose_left.shape == pose_right.shape == (7,)
+    ), f"pose_left.shape: {pose_left.shape}, pose_right.shape: {pose_right.shape}"
+    assert timestamp_left <= timestamp < timestamp_right
+    ratio = (timestamp - timestamp_left) / (timestamp_right - timestamp_left)
+    pos = pose_left[:3] + ratio * (pose_right[:3] - pose_left[:3])
+    rot_left = st.Rotation.from_quat(to_xyzw(pose_left[3:]))
+    rot_right = st.Rotation.from_quat(to_xyzw(pose_right[3:]))
+    rots = st.Rotation.concatenate([rot_left, rot_right])
+    rot = st.Slerp([0, 1], rots)([ratio])
+    return np.concatenate([pos, to_wxyz(rot.as_quat().squeeze())])
+
+
+def get_absolute_pose(
+    init_pose_xyz_wxyz: npt.NDArray[Any],
+    relative_pose_xyz_wxyz: npt.NDArray[Any],
+):
+    """The new pose is in the same frame of reference as the initial pose"""
+    new_pose_xyz_wxyz = np.zeros(7, init_pose_xyz_wxyz.dtype)
+    relative_pos_in_init_frame_as_quat_wxyz = np.zeros(4, init_pose_xyz_wxyz.dtype)
+    relative_pos_in_init_frame_as_quat_wxyz[1:] = relative_pose_xyz_wxyz[:3]
+    init_rot_qinv = qconjugate(init_pose_xyz_wxyz[3:])
+    relative_pos_in_world_frame_as_quat_wxyz = qmult(
+        qmult(init_pose_xyz_wxyz[3:], relative_pos_in_init_frame_as_quat_wxyz),
+        init_rot_qinv,
+    )
+    new_pose_xyz_wxyz[:3] = (
+        init_pose_xyz_wxyz[:3] + relative_pos_in_world_frame_as_quat_wxyz[1:]
+    )
+    quat = qmult(init_pose_xyz_wxyz[3:], relative_pose_xyz_wxyz[3:])
+    if quat[0] < 0:
+        quat = -quat
+    new_pose_xyz_wxyz[3:] = quat
+    return new_pose_xyz_wxyz
+
+
+def get_relative_pose(
+    new_pose_xyz_wxyz: npt.NDArray[Any],
+    init_pose_xyz_wxyz: npt.NDArray[Any],
+):
+    """The two poses are in the same frame of reference"""
+    relative_pose_xyz_wxyz = np.zeros(7, new_pose_xyz_wxyz.dtype)
+    relative_pos_in_world_frame_as_quat_wxyz = np.zeros(4, new_pose_xyz_wxyz.dtype)
+    relative_pos_in_world_frame_as_quat_wxyz[1:] = (
+        new_pose_xyz_wxyz[:3] - init_pose_xyz_wxyz[:3]
+    )
+    init_rot_qinv = qconjugate(init_pose_xyz_wxyz[3:])
+    relative_pose_xyz_wxyz[:3] = qmult(
+        qmult(init_rot_qinv, relative_pos_in_world_frame_as_quat_wxyz),
+        init_pose_xyz_wxyz[3:],
+    )[1:]
+    quat = qmult(init_rot_qinv, new_pose_xyz_wxyz[3:])
+    if quat[0] < 0:
+        quat = -quat
+    relative_pose_xyz_wxyz[3:] = quat
+    return relative_pose_xyz_wxyz
+
+
+def get_relative_poses(
+    poses_xyz_wxyz: npt.NDArray[Any],
+    init_pose_xyz_wxyz: npt.NDArray[Any],
+) -> npt.NDArray[Any]:
+    return np.array(
+        [
+            get_relative_pose(pose_xyz_wxyz, init_pose_xyz_wxyz)
+            for pose_xyz_wxyz in poses_xyz_wxyz
+        ]
+    )
+
+
+def invert_pose(pose_xyz_wxyz: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    qinv = qconjugate(pose_xyz_wxyz[3:])
+    pos_quat_wxyz = np.zeros(4, pose_xyz_wxyz.dtype)
+    pos_quat_wxyz[1:] = pose_xyz_wxyz[:3]
+    rotated_pos = qmult(
+        qmult(qinv, pos_quat_wxyz),
+        pose_xyz_wxyz[3:],
+    )
+    inverted_pose = np.zeros(7, pose_xyz_wxyz.dtype)
+    inverted_pose[:3] = -rotated_pos[1:]
+    if qinv[0] < 0:
+        qinv = -qinv
+    inverted_pose[3:] = qinv
+    return inverted_pose
+
+
+
+
+def normalize(vec: npt.NDArray[Any], eps: float = 1e-12) -> npt.NDArray[Any]:
+    norm: npt.NDArray[Any] = np.linalg.norm(vec, axis=-1)
+    norm = np.maximum(norm, eps)
+    out: npt.NDArray[Any] = (vec.T / norm).T
+    return out
+
+def quat_wxyz_to_rot_6d(quat_wxyz: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    """
+    Convert a quaternion to a 6D representation: the first two rows of the corresponding rotation matrix.
+    https://arxiv.org/pdf/1812.07035
+    quat_wxyz: (4, )
+    return: (6, )
+    """
+    assert quat_wxyz.shape == (4,)
+    w, x, y, z = quat_wxyz[0], quat_wxyz[1], quat_wxyz[2], quat_wxyz[3]
+
+    R = np.array(
+        [
+            [1 - 2 * y * y - 2 * z * z, 2 * x * y - 2 * w * z, 2 * x * z + 2 * w * y],
+            [2 * x * y + 2 * w * z, 1 - 2 * x * x - 2 * z * z, 2 * y * z - 2 * w * x],
+            [2 * x * z - 2 * w * y, 2 * y * z + 2 * w * x, 1 - 2 * x * x - 2 * y * y],
+        ]
+    )
+
+    rot_6d = np.zeros(6)
+    rot_6d[:3] = R[0, :]
+    rot_6d[3:] = R[1, :]
+
+    return rot_6d
+
+
+def rot_6d_to_quat_wxyz(rot_6d: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    """
+    Convert a 6D representation to a quaternion.
+    https://arxiv.org/pdf/1812.07035
+    rot_6d: (6, )
+    return: (4, )
+    """
+
+    assert rot_6d.shape == (6,)
+    a1, a2 = rot_6d[:3], rot_6d[3:]
+    b1 = a1 / np.linalg.norm(a1)
+    b2 = a2 - np.dot(b1, a2) * b1
+    b2 = b2 / np.linalg.norm(b2)
+    b3 = np.cross(b1, b2)
+
+    m = np.zeros((3, 3))
+    m[0, :] = b1
+    m[1, :] = b2
+    m[2, :] = b3
+
+    trace = np.trace(m)
+
+    if trace > 0:
+        s = 0.5 / np.sqrt(trace + 1.0)
+        w = 0.25 / s
+        x = (m[2, 1] - m[1, 2]) * s
+        y = (m[0, 2] - m[2, 0]) * s
+        z = (m[1, 0] - m[0, 1]) * s
+    elif m[0, 0] > m[1, 1] and m[0, 0] > m[2, 2]:
+        s = 2.0 * np.sqrt(1.0 + m[0, 0] - m[1, 1] - m[2, 2])
+        w = (m[2, 1] - m[1, 2]) / s
+        x = 0.25 * s
+        y = (m[0, 1] + m[1, 0]) / s
+        z = (m[0, 2] + m[2, 0]) / s
+    elif m[1, 1] > m[2, 2]:
+        s = 2.0 * np.sqrt(1.0 + m[1, 1] - m[0, 0] - m[2, 2])
+        w = (m[0, 2] - m[2, 0]) / s
+        x = (m[0, 1] + m[1, 0]) / s
+        y = 0.25 * s
+        z = (m[1, 2] + m[2, 1]) / s
+    else:
+        s = 2.0 * np.sqrt(1.0 + m[2, 2] - m[0, 0] - m[1, 1])
+        w = (m[1, 0] - m[0, 1]) / s
+        x = (m[0, 2] + m[2, 0]) / s
+        y = (m[1, 2] + m[2, 1]) / s
+        z = 0.25 * s
+
+    if w < 0:
+        w, x, y, z = -w, -x, -y, -z
+
+    return np.array([w, x, y, z])
+
+
+def quat_wxyz_to_rot_6d_batch(quat_wxyz: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    """
+    input (..., 4)
+    output (..., 6)
+    """
+    assert quat_wxyz.shape[-1] == 4
+    input_shape = quat_wxyz.shape[:-1]
+    quat_wxyz = quat_wxyz.copy().reshape(-1, 4)
+    rot_6d = np.zeros((quat_wxyz.shape[0], 6))
+    for i in range(quat_wxyz.shape[0]):
+        rot_6d[i] = quat_wxyz_to_rot_6d(quat_wxyz[i])
+    return rot_6d.reshape(*input_shape, 6)
+
+
+def rot_6d_to_quat_wxyz_batch(rot_6d: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    """
+    input (..., 6)
+    output (..., 4)
+    """
+    assert rot_6d.shape[-1] == 6
+    input_shape = rot_6d.shape[:-1]
+    rot_6d = rot_6d.copy().reshape(-1, 6)
+    quat_wxyz = np.zeros((rot_6d.shape[0], 4))
+    for i in range(rot_6d.shape[0]):
+        quat_wxyz[i] = rot_6d_to_quat_wxyz(rot_6d[i])
+    return quat_wxyz.reshape(*input_shape, 4)
+
+def rot_6d_to_mat(d6: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    a1, a2 = d6[..., :3], d6[..., 3:]
+    b1: npt.NDArray[Any] = normalize(a1)
+    b2: npt.NDArray[Any] = a2 - np.sum(b1 * a2, axis=-1, keepdims=True) * b1
+    b2: npt.NDArray[Any] = normalize(b2)
+    b3: npt.NDArray[Any] = np.cross(b1, b2, axis=-1)
+    out: npt.NDArray[Any] = np.stack((b1, b2, b3), axis=-2)
+    return out
+
+
+def mat_to_rot_6d(mat: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    batch_dim = mat.shape[:-2]
+    out: npt.NDArray[Any] = mat[..., :2, :].copy().reshape(batch_dim + (6,))
+    return out
+
+
+def pose_9d_to_xyz_wxyz(pose_9d: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    assert pose_9d.shape[-1] == 9, f"pose_9d.shape: {pose_9d.shape}"
+    pos_xyz = pose_9d[..., :3]
+    rot_mqt = rot_6d_to_mat(pose_9d[..., 3:])
+    rot_wxyz = to_wxyz(Rotation.from_matrix(rot_mqt).as_quat())
+    return np.concatenate([pos_xyz, rot_wxyz], axis=-1)
+
+
+def xyz_wxyz_to_pose_9d(pose_xyz_wxyz: npt.NDArray[Any]) -> npt.NDArray[Any]:
+    assert pose_xyz_wxyz.shape[-1] == 7, f"pose_xyz_wxyz.shape: {pose_xyz_wxyz.shape}"
+    pos_xyz = pose_xyz_wxyz[..., :3]
+    rot_mqt = Rotation.from_quat(to_xyzw(pose_xyz_wxyz[..., 3:])).as_matrix()
+    rot_9d = mat_to_rot_6d(rot_mqt)
+    return np.concatenate([pos_xyz, rot_9d], axis=-1)
+
+
+def positive_w(
+    pose_xyz_wxyz: npt.NDArray[np.float64],
+) -> npt.NDArray[np.float64]:
+    if pose_xyz_wxyz[3] < 0.0:
+        pose_xyz_wxyz[3:] = -pose_xyz_wxyz[3:]
+    return pose_xyz_wxyz

robot_utils-0.1.1/src/robot_utils/torch_utils.py

@@ -0,0 +1,195 @@
+import io
+import boto3
+import torch
+import os
+from typing import Any, Callable
+from collections.abc import Iterator
+import torch.nn as nn
+import numpy as np
+
+def torch_load(data_path: str, **kwargs):
+    if data_path.startswith("s3://"):
+        s3 = boto3.client("s3")
+        buffer = io.BytesIO()
+        bucket_name, key = data_path.replace("s3://", "").split("/", 1)
+        s3.download_fileobj(bucket_name, key, buffer)
+        buffer.seek(0)
+        data = torch.load(buffer, **kwargs)
+    else:
+        data = torch.load(data_path, **kwargs)
+    return data
+
+def torch_save(data: Any, path: str | io.BufferedWriter, **kwargs):
+    if isinstance(path, str) and path.startswith("s3://"):
+        s3 = boto3.client("s3")
+        buffer = io.BytesIO()
+        torch.save(data, buffer, **kwargs)
+        buffer.seek(0)
+        bucket_name, key = path.replace("s3://", "").split("/", 1)
+        s3.upload_fileobj(buffer, bucket_name, key)
+    else:
+        if isinstance(path, str) and not os.path.exists(os.path.dirname(path)):
+            os.makedirs(os.path.dirname(path))
+        torch.save(data, path, **kwargs)
+
+def filter_params(
+    named_params: Iterator[tuple[str, nn.Parameter]],
+    keywords: list[str] | None,
+    requires_grad: bool | None = None,
+) -> Iterator[tuple[str, nn.Parameter]]:
+    if keywords is not None:
+        keywords_used = np.zeros(len(keywords), dtype=bool)
+    else:
+        keywords_used = None
+
+    for i, (name, param) in enumerate(named_params):
+        if keywords is None or any(keyword in name for keyword in keywords):
+            if keywords_used is not None and keywords is not None:
+                keywords_used = keywords_used | np.array(
+                    [keyword in name for keyword in keywords]
+                )
+
+            if requires_grad is not None:
+                if param.requires_grad == requires_grad:
+                    yield name, param
+            else:
+                yield name, param
+    if keywords_used is not None and keywords is not None:
+        for i, keyword in enumerate(keywords):
+            if not keywords_used[i]:
+                raise ValueError(f"Keyword {keyword} not used! Please check the names")
+
+
+def exclude_params(
+    named_params: Iterator[tuple[str, nn.Parameter]],
+    keywords: list[str],
+    requires_grad: bool | None = None,
+) -> Iterator[tuple[str, nn.Parameter]]:
+    for name, param in named_params:
+        if not any(keyword in name for keyword in keywords):
+            if requires_grad is not None:
+                if param.requires_grad == requires_grad:
+                    yield name, param
+            else:
+                yield name, param
+
+
+def params(named_params: Iterator[tuple[str, nn.Parameter]]) -> Iterator[nn.Parameter]:
+    for _, param in named_params:
+        yield param
+
+
+def aggregate_batch(
+    batch: list[Any], aggregate_fn: Callable[[list[Any]], Any], merge_none: bool = True
+) -> Any:
+    """
+    Custom collate function to concatenate nested tensors/ndarray/float along a specified axis.
+    If merge_none is True, the field that has None values will be merged into a single None value. Otherwise will return a list of None values.
+    Popular choices of aggregate_fn:
+        - partial(torch.cat, dim=existing_dim), if you want to concatenate along an existing dimension
+        - partial(torch.stack, dim=new_dim), if you want to stack to a new dimension
+
+    Args:
+        batch (List[Any]): A list of samples from the dataset.
+        aggregate_fn (Callable[[list[Any]], Any]): The function to aggregate the tensors/ndarray/float.
+
+    Returns:
+        Any: The concatenated batch.
+    """
+    if len(batch) == 0:
+        return batch
+    elem = batch[0]
+    if (
+        isinstance(elem, torch.Tensor)
+        or isinstance(elem, np.ndarray)
+        or isinstance(elem, float)
+    ):
+        return aggregate_fn(batch)
+    elif isinstance(elem, dict):
+        return {
+            key: aggregate_batch([d[key] for d in batch], aggregate_fn)
+            for key in elem.keys()
+        }
+    elif isinstance(elem, list):
+        return [aggregate_batch(samples, aggregate_fn) for samples in zip(*batch)]
+    elif elem is None:
+        if merge_none:
+            return None
+    else:
+        return batch
+
+
+def split_batch(
+    batch: Any, split_fn: Callable[[torch.Tensor], tuple[torch.Tensor, ...]]
+) -> Iterator[Any]:
+    """
+    Split a batch into multiple batches along a specified dimension.
+    Popular choices of split_fn:
+        - partial(torch.split, dim=existing_dim), if you want to split along an existing dimension
+        - partial(torch.unbind, dim=diminishing_dim), if you want to split and diminish a dimension
+
+    Args:
+        batch (Any): Should be a nested dict or a nested list, where all the elements are tensors.
+        split_fn (Callable[[torch.Tensor], tuple[torch.Tensor, ...]]): The function to split the batch.
+
+    Returns:
+        Iterator[Any]: An iterator over the split batches.
+    """
+    if isinstance(batch, torch.Tensor):
+        yield from split_fn(batch)
+
+    elif isinstance(batch, dict):
+        for values in zip(*[split_batch(v, split_fn) for v in batch.values()]):
+            yield {k: v for k, v in zip(batch.keys(), values)}
+
+    elif isinstance(batch, list):
+        for values in zip(*[split_batch(v, split_fn) for v in batch]):
+            yield [v for v in values]
+
+    else:
+        raise ValueError(f"Invalid batch type: {type(batch)}")
+
+
+def to_cpu(obj: Any) -> Any:
+    if isinstance(obj, torch.Tensor):
+        return obj.cpu()
+    elif isinstance(obj, dict):
+        return {key: to_cpu(value) for key, value in obj.items()}
+    elif isinstance(obj, list):
+        return [to_cpu(item) for item in obj]
+    else:
+        return obj
+
+process_group_initialized = False
+
+world_size = int(os.environ.get("WORLD_SIZE", 1))
+
+def init_process_group():
+    global process_group_initialized
+    if not process_group_initialized and world_size > 1:
+        print(f"Environment variable WORLD_SIZE={os.environ.get('WORLD_SIZE')}, RANK={os.environ.get('RANK')}. Using distributed training.")
+        torch.distributed.init_process_group(backend="nccl")
+        process_group_initialized = True
+
+def is_main_process():
+    init_process_group()
+    if torch.distributed.is_initialized():
+        return int(os.environ.get("RANK", 0)) == 0
+    else:
+        return True
+
+def wait_for_main_process():
+    init_process_group()
+    if torch.distributed.is_initialized():
+        # BUG: This will lead to memory leak on GPU 0. Not sure which part of the code is causing this.
+        # Do not use this function before accelerator is initialized.
+        torch.distributed.barrier()
+    else:
+        print("Not using distributed training. No need to wait for main process.")
+    
+def num_processes():
+    init_process_group()
+    if torch.distributed.is_initialized():
+        return int(os.environ.get("WORLD_SIZE", 1))
+    else:
+        return 1

robot_utils-0.1.1/src/robot_utils/video_utils.py

@@ -0,0 +1,20 @@
+from typing import Any
+import numpy as np
+import numpy.typing as npt
+import imageio
+
+def save_np_array_as_video(
+    rollout_images: list[npt.NDArray[Any]] | npt.NDArray[np.uint8],
+    video_path: str,
+    fps: int = 30,
+):
+    """Saves an MP4 replay of an episode."""
+
+    video_writer = imageio.get_writer(video_path, fps=fps)
+    for img in rollout_images:
+        if img.dtype != np.uint8:
+            img = (img * 255).astype(np.uint8)
+        video_writer.append_data(img)
+    video_writer.close()
+    print(f"Saved rollout MP4 at path {video_path}")
+

robot_utils-0.1.1/src/robot_utils.egg-info/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: robot-utils
+Version: 0.1.1
+Summary: Frequently used utility functions for robot learning research
+Author-email: Yihuai Gao <davidgao1013@gmail.com>
+Project-URL: Homepage, https://github.com/yihuai-gao/robot-utils
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent

robot_utils-0.1.1/src/robot_utils.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+pyproject.toml
+src/robot_utils/__init__.py
+src/robot_utils/config_utils.py
+src/robot_utils/data_utils.py
+src/robot_utils/logging_utils.py
+src/robot_utils/pose_utils.py
+src/robot_utils/torch_utils.py
+src/robot_utils/video_utils.py
+src/robot_utils.egg-info/PKG-INFO
+src/robot_utils.egg-info/SOURCES.txt
+src/robot_utils.egg-info/dependency_links.txt
+src/robot_utils.egg-info/top_level.txt

robot_utils-0.1.1/src/robot_utils.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

robot_utils-0.1.1/src/robot_utils.egg-info/top_level.txt

@@ -0,0 +1 @@
+robot_utils