ezmujoco 0.1.0__tar.gz

ezmujoco-0.1.0/PKG-INFO

@@ -0,0 +1,23 @@
+Metadata-Version: 2.3
+Name: ezmujoco
+Version: 0.1.0
+Summary: Add your description here
+Requires-Dist: ezpose>=0.1.16
+Requires-Dist: imageio>=2.37.2
+Requires-Dist: lxml>=6.0.2
+Requires-Dist: mujoco>=3.3.7
+Requires-Dist: yourdfpy>=0.0.58
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# ezmujoco
+
+`ezmujoco` is a lightweight utility library for building and manipulating MuJoCo scenes in Python.
+It provides small, practical abstractions around `Scene`, `World`, and `Robot` so you can compose objects from XML, manage robot models, and work with poses, mocap bodies, cameras, and simple simulation workflows with less boilerplate.
+
+## Development
+
+```bash
+uv sync --dev
+uv run pytest -q
+```

ezmujoco-0.1.0/README.md

@@ -0,0 +1,11 @@
+# ezmujoco
+
+`ezmujoco` is a lightweight utility library for building and manipulating MuJoCo scenes in Python.
+It provides small, practical abstractions around `Scene`, `World`, and `Robot` so you can compose objects from XML, manage robot models, and work with poses, mocap bodies, cameras, and simple simulation workflows with less boilerplate.
+
+## Development
+
+```bash
+uv sync --dev
+uv run pytest -q
+```

ezmujoco-0.1.0/pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "ezmujoco"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "ezpose>=0.1.16",
+    "imageio>=2.37.2",
+    "lxml>=6.0.2",
+    "mujoco>=3.3.7",
+    "yourdfpy>=0.0.58",
+]
+
+[dependency-groups]
+dev = [
+    "pytest>=8.4.2",
+]
+
+[build-system]
+requires = ["uv_build>=0.9.6,<0.10.0"]
+build-backend = "uv_build"

ezmujoco-0.1.0/src/ezmujoco/__init__.py

@@ -0,0 +1,14 @@
+from .robot import Robot
+from .scene import Scene, SpecInfo
+from .world import World, Body, Camera, MocappedBody, VideoMaker
+
+__all__ = [
+    "Robot",
+    "Scene",
+    "SpecInfo",
+    "World",
+    "Body",
+    "Camera",
+    "MocappedBody",
+    "VideoMaker",
+]

ezmujoco-0.1.0/src/ezmujoco/hand_util.py

@@ -0,0 +1,303 @@
+import numpy as np
+from . import Robot, Body
+from ezpose import SE3, SO3
+
+# class Hand(Robot):
+#     def __init__(self, palm_pose=None, **kwargs):
+#         super().__init__(**kwargs)
+#         self.pose_base_palm = palm_pose
+    
+#     def set_palm_pose(self, palm_pose: SE3):
+#         pose_base = palm_pose @ self.pose_base_palm.inv()
+#         self.set_base_pose(pose_base)
+
+#     def get_palm_pose(self):
+#         return self.get_base_pose() @ self.pose_base_palm
+    
+#     def ctrl_palm_pose(self, palm_pose):
+#         pose_base = palm_pose @ self.pose_base_palm.inv()
+#         self.ctrl_base_pose(pose_base)
+    
+#     def converge_finger(self, q):
+#         self.set_pos_ctrl_target(q)
+#         self.world.wait_to_stabilize(timeout=2.0)
+#         return self.get_joint_angles()
+    
+#     def move_xyz(
+#         self,
+#         delta_xyz,
+#         substep=10,
+#         duration=1.0,
+#         rest=0.1,
+#         grasp_target: Body = None,
+#     ):
+#         delta_xyz = np.array(delta_xyz)
+#         curr_pose = self.get_palm_pose()
+#         self.ctrl_palm_pose(curr_pose)
+
+#         n_steps = int(duration / self.world.dt)
+#         n_substeps = n_steps // substep
+#         assert n_steps > n_substeps
+
+#         final_pose = SE3(trans=delta_xyz) @ curr_pose
+
+#         ratio_grid = np.linspace(0, 1, n_substeps, endpoint=True)
+#         pose_grid = [curr_pose.interpolate(final_pose, r) for r in ratio_grid]
+
+#         for i in range(n_substeps):
+#             if grasp_target is not None:
+#                 if not self.is_in_collision(grasp_target):
+#                     break
+#             self.ctrl_palm_pose(pose_grid[i])
+#             for _ in range(substep):
+#                 self.world.step()
+
+#         for _ in range(int(rest / self.world.dt)):
+#             if grasp_target is not None:
+#                 if not self.is_in_collision(grasp_target):
+#                     break
+#             self.world.step()
+    
+#     def move_pose(
+#         self,
+#         final_pose: SE3,
+#         substep=10,
+#         duration=1.0,
+#         rest=0.1,
+#         grasp_target: Body = None,
+#         abort_on_contact: bool = False
+#     ):
+#         n_steps = int(duration / self.world.dt)
+#         n_substeps = n_steps // substep
+#         assert n_steps > n_substeps
+
+#         curr_pose: SE3 = self.get_palm_pose()
+#         self.ctrl_palm_pose(curr_pose)
+
+#         ratio_grid = np.linspace(0, 1, n_substeps, endpoint=True)
+#         pose_grid = []
+#         for r in ratio_grid:
+#             trans = curr_pose.trans + (final_pose.trans - curr_pose.trans) * r
+#             rot = curr_pose.rot.interpolate(final_pose.rot, r)
+#             pose_grid.append(SE3(trans=trans, rot=rot))
+
+#         for i in range(n_substeps):
+#             if grasp_target is not None:
+#                 if not self.is_in_collision(grasp_target):
+#                     break
+#             if abort_on_contact:
+#                 if self.is_in_collision(exclude_list=self.bodies):
+#                     break
+                
+#             self.ctrl_palm_pose(pose_grid[i])
+#             for _ in range(substep):
+#                 self.world.step()
+
+#         for _ in range(int(rest / self.world.dt)):
+#             if grasp_target is not None:
+#                 if not self.is_in_collision(grasp_target):
+#                     break
+#             self.world.step()
+
+class Hand(Robot):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.pose_base_palm: SE3 = None
+        self.span_idx: np.ndarray = None
+        self.shrink_idx: np.ndarray = None
+        self.mimic_idx: np.ndarray = None
+        self.set_color(alpha=1.)
+    
+    @classmethod
+    def parse_from_xml(
+        cls,
+        robot_name,
+        world,
+        xml_path,
+        palm_pose: SE3,
+        joint_order: list[str] = None,
+        joint_order_map: dict[str, str] = None,
+    ):
+        cls = super().parse_from_xml(
+            robot_name,
+            world,
+            xml_path,
+            joint_order=joint_order,
+            joint_order_map=joint_order_map,
+        )
+        cls.pose_base_palm = palm_pose
+        return cls
+
+    def set_color(self, alpha=1.):
+        for body in self.bodies:
+            if "thumb" in body.name:
+                rgb = [0, 0, 255] # [5, 84, 242] # b
+            elif "index" in body.name:
+                rgb = [0, 255, 0] # [10, 191, 4] # g
+            elif "middle" in body.name:
+                rgb = [255, 0, 0] # [242, 82, 46] # r
+            elif "ring" in body.name:
+                rgb = [255, 150, 0] #[242, 159, 5] # y
+            else:
+                rgb = [55] * 3          
+            rgba = [*np.array(rgb)/255, alpha]
+            body.set_color(rgba=rgba)
+
+    def register_palm_pose(self, palm_pose: SE3):
+        self.pose_base_palm = palm_pose
+
+    def set_palm_pose(self, palm_pose: SE3):
+        assert self.pose_base_palm is not None, "Please register palm pose first"
+        pose_base = palm_pose @ self.pose_base_palm.inv()
+        self.set_base_pose(pose_base)
+        
+        self.bodies[0].set_pose(pose_base)
+        if self.mocap_body is not None:
+            self.ctrl_base_pose(pose_base)
+        
+
+    def get_palm_pose(self):
+        return self.get_base_pose() @ self.pose_base_palm
+
+    def ctrl_palm_pose(self, palm_pose):
+        pose_base = palm_pose @ self.pose_base_palm.inv()
+        self.ctrl_base_pose(pose_base)
+
+    def converge_finger(self, q, duration=1.0, wait=1.0):
+        q_init = self.get_joint_angles()
+        dt = self.world.dt
+        steps = int(duration / dt)
+        for i in range(steps):
+            ratio = i / steps
+            q_int = ratio * q + (1-ratio) * q_init
+            self.set_pos_ctrl_target(q_int)
+            self.world.step()
+        self.set_pos_ctrl_target(q)
+        self.world.wait_to_stabilize(timeout=2.0)
+        for _ in range(int(wait / dt)): # wait for 100 ms
+            self.world.step()
+        return self.get_joint_angles()
+
+    def set_ghost(self, value: bool):
+        for body in self.bodies:
+            body.set_ghost(value)
+        if value:
+            self.set_color(0.)
+        else:
+            self.set_color(1.)
+
+
+
+    def move_xyz(
+        self,
+        delta_xyz,
+        substep=10,
+        duration=1.0,
+        rest=0.1,
+        grasp_target: Body = None,
+    ):
+        delta_xyz = np.array(delta_xyz)
+        curr_pose = self.get_palm_pose()
+        self.ctrl_palm_pose(curr_pose)
+
+        n_steps = int(duration / self.world.dt)
+        n_substeps = n_steps // substep
+        assert n_steps > n_substeps
+
+        final_pose = SE3(trans=delta_xyz) @ curr_pose
+
+        ratio_grid = np.linspace(0, 1, n_substeps, endpoint=True)
+        pose_grid = [curr_pose.interpolate(final_pose, r) for r in ratio_grid]
+
+        for i in range(n_substeps):
+            if grasp_target is not None:
+                if not self.is_in_collision(grasp_target):
+                    break
+            self.ctrl_palm_pose(pose_grid[i])
+            for _ in range(substep):
+                self.world.step()
+
+        for _ in range(int(rest / self.world.dt)):
+            if grasp_target is not None:
+                if not self.is_in_collision(grasp_target):
+                    break
+            self.world.step()
+
+    def move_rpy(
+        self,
+        delta: float,
+        substep=10,
+        duration=1.0,
+        rest=0.1,
+        rpy="roll",
+        grasp_target: Body = None,
+    ):
+        curr_pose = self.get_palm_pose()
+        self.ctrl_palm_pose(curr_pose)
+
+        n_steps = int(duration / self.world.dt)
+        n_substeps = n_steps // substep
+        assert n_steps > n_substeps
+
+        dist_step = delta / n_substeps
+        for _ in range(n_substeps):
+            if grasp_target is not None:
+                if not self.is_in_collision(grasp_target):
+                    break
+            if rpy == "roll":
+                ypr_delta = [0, 0, dist_step]
+            elif rpy == "pitch":
+                ypr_delta = [0, dist_step, 0]
+            elif rpy == "yaw":
+                ypr_delta = [dist_step, 0, 0]
+            curr_pose.rot = curr_pose.rot @ SO3.from_euler("ZYX", ypr_delta)
+            self.ctrl_palm_pose(curr_pose)
+            for _ in range(substep):
+                self.world.step()
+
+        for _ in range(int(rest / self.world.dt)):
+            if grasp_target is not None:
+                if not self.is_in_collision(grasp_target):
+                    break
+            self.world.step()
+
+    def move_pose(
+        self,
+        final_pose: SE3,
+        substep=10,
+        duration=1.0,
+        rest=0.1,
+        grasp_target: Body = None,
+        abort_on_contact: bool = False
+    ):
+        n_steps = int(duration / self.world.dt)
+        n_substeps = n_steps // substep
+        assert n_steps > n_substeps
+
+        curr_pose: SE3 = self.get_palm_pose()
+        self.ctrl_palm_pose(curr_pose)
+
+        ratio_grid = np.linspace(0, 1, n_substeps, endpoint=True)
+        pose_grid = []
+        for r in ratio_grid:
+            trans = curr_pose.trans + (final_pose.trans - curr_pose.trans) * r
+            rot = curr_pose.rot.interpolate(final_pose.rot, r)
+            pose_grid.append(SE3(trans=trans, rot=rot))
+
+        for i in range(n_substeps):
+            if grasp_target is not None:
+                if not self.is_in_collision(grasp_target):
+                    break
+            if abort_on_contact:
+                if self.is_in_collision(exclude_list=self.bodies):
+                    break
+                
+            self.ctrl_palm_pose(pose_grid[i])
+            for _ in range(substep):
+                self.world.step()
+
+        for _ in range(int(rest / self.world.dt)):
+            if grasp_target is not None:
+                if not self.is_in_collision(grasp_target):
+                    break
+            self.world.step()