egogym 0.1.0__py3-none-any.whl

egogym/components/__init__.py

@@ -0,0 +1 @@
+from .object import Object, ArticulableObject

egogym/components/object.py

@@ -0,0 +1,94 @@
+import mujoco
+from scipy.spatial.transform import Rotation as R
+import numpy as np
+
+from egogym.utils import get_pose, include_in_scene, add_xml_to_scene
+
+
+class Object:
+    def __init__(self, name, path=None, func=None, np_random=None):
+        self.name = name
+        self.english_name = self.name.rsplit("_", 1)[0].replace("_", " ")
+        self.path = path
+        self.last_set_pose = None
+        self.object_xml = None
+        self.handle_type = None
+        self.object_geoms_ids = None
+        self.np_random = np_random
+        if func is not None:
+            self.object_xml = func(
+                np_random=np_random
+            )
+        self.body_name = f"{self.name}_object"
+
+    def get_pose(self, data):
+        return get_pose(data, self.body_name, "body")
+
+    def get_geom_ids(self, model):
+        if self.object_geoms_ids is not None:
+            return self.object_geoms_ids
+        object_geoms_ids = []
+        for i in range(model.ngeom):
+            geom_bodyid = model.geom_bodyid[i]
+            body_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, geom_bodyid)
+            if body_name is not None and f"{self.name}_object" in body_name:
+                object_geoms_ids.append(i)
+        self.object_geoms_ids = object_geoms_ids
+        return object_geoms_ids
+
+    def get_bottom_pose(self, data):
+        return get_pose(data, f"{self.name}_bottom_site", "site")
+
+    def get_center_to_bottom_z_distance(self, data):
+        bottom_z = self.get_bottom_pose(data)[2, 3]
+        center_z = self.get_pose(data)[2, 3]
+        return center_z - bottom_z
+
+    def set_pose(self, model, data, pose):
+        body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, self.body_name)
+        model.body_pos[body_id] = pose[0:3, 3]
+        model.body_quat[body_id] = R.from_matrix(pose[0:3, 0:3]).as_quat()
+        self.last_set_pose = pose.copy()
+
+    def add_to_scene_xml(self, scene_xml):
+        if self.object_xml is not None:
+            return add_xml_to_scene(scene_xml, self.object_xml)
+        elif self.path is not None:
+            return include_in_scene(scene_xml, self.path)
+        else:
+            raise ValueError("Either func or path must be provided")
+
+
+class ArticulableObject(Object):
+    def __init__(self, name, path=None, func=None, np_random=None):
+        super().__init__(name, path, func, np_random)
+        self.joint_range = None
+        self.initial_joint = None
+    
+    def get_handle_pose(self, data):
+        return get_pose(data, "handle", "geom")
+    
+    def get_jotinpos(self, data):
+        return data.joint("task_joint").qpos[0]
+    
+    def get_perecenttage_opened(self, model, data):
+        if self.joint_range is None:
+            self.joint_range = model.jnt_range[mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "task_joint")]
+            self.initial_joint = self.get_jotinpos(data)
+        current_task_joint = self.get_jotinpos(data)
+        normalized_task_joint = abs(current_task_joint - self.initial_joint) / (
+            abs(self.joint_range[1] - self.joint_range[0]) + 1e-8
+        )
+        return normalized_task_joint
+    
+    def open(self, model, data):
+        if self.joint_range is None:
+            self.joint_range = model.jnt_range[mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "task_joint")]
+            self.initial_joint = self.get_jotinpos(data)
+        data.joint("task_joint").qpos = (
+            self.joint_range[1] * self.np_random.uniform(0.5, 0.9)
+            if self.joint_range[0] == 0
+            else self.joint_range[0] * self.np_random.uniform(0.5, 0.9)
+        )
+        mujoco.mj_step(model, data, nstep=100)
+        

egogym/egogym.py

@@ -0,0 +1,106 @@
+from gymnasium.utils import seeding
+import gymnasium as gym
+from gymnasium.spaces import Box, Dict
+import mujoco
+import numpy as np
+import cv2
+
+from egogym.embodiments import ROBOT_REGISTRY
+
+class Egogym(gym.Env):
+    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 30}
+
+    def __init__(self, robot="rum", action_space="delta", render_mode=None, render_size=(960,720), seed=None):
+        super().__init__()
+        self.model = None
+        self.data = None
+        self.step_idx = 0
+        self.render_mode = render_mode
+        self.np_random, _ = seeding.np_random(seed)
+    
+        self.render_width = render_size[0]
+        self.render_height = render_size[1]
+        
+        self.robot_type = robot
+        self.action_space_type = action_space
+        
+        if self.robot_type not in ROBOT_REGISTRY:
+            raise ValueError(f"Unknown robot type '{self.robot_type}'. Available robots: {list(ROBOT_REGISTRY.keys())}")
+        
+        robot_class = ROBOT_REGISTRY[self.robot_type]
+        self.robot = robot_class(action_space=self.action_space_type, np_random=self.np_random)
+        
+        obs_spaces = {
+            "rgb_ego": Box(low=0, high=255, shape=(self.render_height, self.render_width, 3), dtype=np.uint8),
+            "rgb_exo": Box(low=0, high=255, shape=(self.render_height, self.render_width, 3), dtype=np.uint8),
+            "tcp_pose": Box(low=-np.inf, high=np.inf, shape=(16,), dtype=np.float32),
+            "camera_pose": Box(low=-np.inf, high=np.inf, shape=(16,), dtype=np.float32),
+            "grasp": Box(low=-np.inf, high=np.inf, shape=(1,), dtype=np.float32),
+        }
+        obs_spaces.update(self.robot.additional_obs_spaces())
+        self.observation_space = Dict(obs_spaces)
+        self.action_space = Box(low=-np.inf, high=np.inf, shape=self.robot.action_space_shape(), dtype=np.float32)
+        self._mujoco_step_counter = 0
+        self._render_freq = 0
+        self._render_callback = None
+        self._should_render = self.render_mode == 'human'
+
+    def setup_mujoco(self, scene_xml: str):
+        if hasattr(self, "model") and self.model is not None:
+            del self.model
+            del self.data
+        self.model = mujoco.MjModel.from_xml_string(scene_xml)
+        self.data = mujoco.MjData(self.model)
+        self.robot.setup_mj(self.model, self.data)
+        self.robot.setup_renderers(self.render_width, self.render_height)
+
+    def env_step(self, nsteps=1):
+        for _ in range(nsteps):
+            mujoco.mj_step(self.model, self.data)
+            self._mujoco_step_counter += 1
+            if self._render_freq > 0 and self._render_callback is not None and self._mujoco_step_counter >= self._render_freq:
+                self._render_callback()
+                self._mujoco_step_counter = 0
+
+    def get_base_obs(self):
+        obs = {
+            "rgb_ego": self.robot.get_camera_view(self.robot.camera_names[0]),
+            "rgb_exo": self.robot.get_camera_view(self.robot.camera_names[-1]),
+            "tcp_pose": self.robot.get_tcp_pose().reshape(16).astype(np.float32),
+            "camera_pose": self.robot.get_camera_pose().reshape(16).astype(np.float32),
+            "grasp": np.array([self.robot.get_grasp()], dtype=np.float32),
+        }
+        obs.update(self.robot.get_additional_obs())
+        return obs
+
+    def render(self):
+        rgb_ego = self.robot.get_camera_view(self.robot.camera_names[0])
+        rgb_exo = self.robot.get_camera_view(self.robot.camera_names[1])
+        rgb_stack = np.concatenate((rgb_ego, rgb_exo), axis=1)
+        rgb_stack = cv2.cvtColor(rgb_stack, cv2.COLOR_RGB2BGR)
+        cv2.imshow("Egogym", rgb_stack)
+        cv2.waitKey(1)
+
+    def close(self):
+        if self.render_mode == "human":
+            cv2.destroyAllWindows()
+
+    def step(self, action):
+        if action is None or np.isnan(action).any():
+            return self.get_obs(), 0.0, False, False, {}
+        self.robot.step(action, step_fn=self.env_step)
+        obs = self.get_obs()
+        reward = self.compute_reward()
+        terminated = False
+        truncated = False
+        return obs, reward, terminated, truncated, {}
+    
+    def enable_sleeping_islands(self):
+        try:
+            self.model.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_SLEEP
+        except AttributeError:
+            pass
+    
+    def settle_env(self, nsteps=200):
+        if hasattr(self.robot, 'free_joint_id') and self.robot.free_joint_id is not None: #TODO: check if its floating_gripper instead
+            self.data.qpos[:] = np.where((np.arange(len(self.data.qpos)) >= self.robot.free_joint_id) & (np.arange(len(self.data.qpos)) < self.robot.free_joint_id + 6), self.data.qpos, 0.0)

egogym/embodiments/__init__.py

@@ -0,0 +1,10 @@
+from .robot import Robot
+from .grippers import FloatingGripper, RUMGripper
+from .arms import Arm, DroidArm
+
+ROBOT_REGISTRY = {
+    "rum": RUMGripper,
+    "droid": DroidArm,
+}
+
+__all__ = ["Robot", "FloatingGripper", "RUMGripper", "Arm", "DroidArm", "ROBOT_REGISTRY"]

egogym/embodiments/arms/__init__.py

@@ -0,0 +1,4 @@
+from .arm import Arm
+from .droid import DroidArm
+
+__all__ = ["Arm", "DroidArm"]

egogym/embodiments/arms/arm.py

@@ -0,0 +1,65 @@
+from egogym.embodiments.robot import Robot, requires_mj
+from gymnasium.spaces import Box
+import numpy as np
+
+
+class Arm(Robot):
+    
+    def __init__(self, control_steps=2000, grasping_steps=1000, action_space="delta", np_random=None):
+        super().__init__(control_steps=control_steps, grasping_steps=grasping_steps, action_space=action_space, np_random=np_random)
+        self.joint_names = []
+        self.qpos_ids = None
+        self.act_ids = None
+        self.gripper_act_id = None
+    
+    def additional_obs_spaces(self):
+        num_joints = len(self.joint_names)
+        return {
+            "joint_positions": Box(low=-np.inf, high=np.inf, shape=(num_joints,), dtype=np.float32)
+        }
+    
+    def action_space_shape(self):
+        return (len(self.joint_names) + 1,)
+    
+    @requires_mj
+    def get_additional_obs(self):
+        return {
+            "joint_positions": self.get_joints().astype(np.float32)
+        }
+    
+    def setup_mj(self, model, data):
+        super().setup_mj(model, data)
+        if self.joint_names:
+            self.qpos_ids = np.array([self.model.joint(name).qposadr[0] for name in self.joint_names])
+            self.act_ids = np.array([self.model.actuator(name).id for name in self.joint_names])
+        if hasattr(self, 'gripper_actuator_name'):
+            self.gripper_act_id = self.model.actuator(self.gripper_actuator_name).id
+    
+    @requires_mj
+    def get_joints(self):
+        return self.data.qpos[self.qpos_ids].copy()
+    
+    @requires_mj
+    def set_joints(self, joint_positions):
+        self.data.ctrl[self.act_ids] = joint_positions
+    
+    @requires_mj
+    def step(self, action, step_fn):
+
+        if len(action.shape) == 2:
+            action = action[0]
+        
+        num_joints = len(self.joint_names)
+        desired_joints = action[:num_joints]
+        grip_action = action[num_joints] if len(action) > num_joints else 0.0
+        
+        if self.action_space == "delta":
+            target_joints = self.get_joints() + desired_joints
+        else:
+            target_joints = desired_joints
+        
+
+        self.set_joints(target_joints)
+        self.data.ctrl[self.gripper_act_id] = 255.0 * np.clip(grip_action, 0.0, 1.0)
+        
+        step_fn(self.control_steps)

egogym/embodiments/arms/droid.py

@@ -0,0 +1,49 @@
+from egogym.embodiments.arms.arm import Arm
+from egogym.embodiments.robot import requires_mj
+from egogym.utils import get_pose
+import numpy as np
+
+
+class DroidArm(Arm):    
+    def __init__(self, control_steps=600, grasping_steps=100, action_space="delta", np_random=None):
+        super().__init__(control_steps=control_steps, grasping_steps=grasping_steps, action_space=action_space, np_random=np_random)
+        self.name = "droid"
+        self.initial_q = np.array([0, -1/5 * np.pi, 0, -4/5 * np.pi, 0, 3/5 * np.pi, 0.0])
+        self.joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"]
+        self.gripper_actuator_name = "fingers_actuator"
+    
+    @requires_mj
+    def prepare(self, desired_pose, step_fn=None):
+
+        desired_pose[:3,3] += np.array([0.0, 0.6, 0.075])
+        if self.add_noise_init:
+            position_noise = self.np_random.uniform(-0.03, 0.03, size=3)
+            desired_pose[:3, 3] += position_noise
+
+        self.data.ctrl[self.act_ids] = self.initial_q
+        step_fn(500)
+        gripper_position = self.get_tcp_pose()[:3, 3]
+        link_position = get_pose(self.data, "link0", "body")[:3, 3]
+        delta_position = gripper_position - link_position
+        desired_gripper_position = desired_pose[:3, 3]
+        desired_link_position = desired_gripper_position - delta_position
+        self.model.body_pos[self.model.body("link0").id] = desired_link_position
+        step_fn(100)
+
+    @requires_mj
+    def get_grasp(self):
+        val = self.data.qpos[self.model.joint("right_driver_joint").qposadr[0]] / 0.8102
+        if val < 0.01:
+            val = 0.0
+        val = np.array([val], dtype=np.float32).reshape(1)
+        current = np.array([(np.clip(val, 0.0, 1.0))], dtype=np.float32).reshape(1)[0]
+        return current
+    
+    @requires_mj
+    def set_grasp(self, grip_action):
+        desired_grasp = 255.0 * grip_action
+        self.data.ctrl[self.model.actuator("fingers_actuator").id] = desired_grasp.item()
+        if desired_grasp != self.last_grasp:
+            self.last_grasp = desired_grasp
+            return 1
+        return 0

egogym/embodiments/grippers/__init__.py

@@ -0,0 +1,4 @@
+from .floating_gripper import FloatingGripper
+from .rum import RUMGripper
+
+__all__ = ["FloatingGripper", "RUMGripper"]

egogym/embodiments/grippers/floating_gripper.py

@@ -0,0 +1,58 @@
+from scipy.spatial.transform import Rotation as R
+import numpy as np
+
+from egogym.utils import get_pose
+from egogym.embodiments.robot import Robot, requires_mj
+
+
+class FloatingGripper(Robot):    
+    def __init__(self, control_steps=50, grasping_steps=1, action_space="delta", np_random=None):
+        super().__init__(control_steps=control_steps, grasping_steps=grasping_steps, action_space=action_space, np_random=np_random)
+        self.last_grasp = 0.0
+    
+    def action_space_shape(self):
+        return (17,)
+
+    def setup_mj(self, model, data):
+        super().setup_mj(model, data)
+        self.last_grasp = 0.0
+        self.free_joint_id = model.joint("free_joint").qposadr[0]
+    
+    @requires_mj
+    def prepare(self, desired_pose, step_fn=None):
+        if self.add_noise_init:
+            position_noise = self.np_random.uniform(-0.03, 0.03, size=3)
+            desired_pose[:3, 3] += position_noise
+            euler_angles = R.from_matrix(desired_pose[:3, :3]).as_euler("xyz")
+            euler_angles[0] += self.np_random.uniform(-0.1, 0.1)
+            euler_angles[1] += self.np_random.uniform(-0.05, 0.05)
+            euler_angles[2] += self.np_random.uniform(-0.1, 0.1)
+            desired_pose[:3, :3] = R.from_euler("xyz", euler_angles).as_matrix()
+        self.last_teleop_pose = desired_pose
+        self.data.mocap_pos[0] = desired_pose[:3, 3]
+        self.data.mocap_quat[0] = np.roll(R.from_matrix(desired_pose[:3, :3]).as_quat(), 1)
+        step_fn(self.control_steps)
+
+    @requires_mj
+    def move(self, desired_pose):
+        self.data.mocap_pos[0] = desired_pose[:3, 3]
+        self.data.mocap_quat[0] = np.roll(R.from_matrix(desired_pose[:3, :3]).as_quat(), 1)
+
+    @requires_mj
+    def step(self, action, step_fn):
+        if len(action.shape) == 2:
+            action = action[0]
+
+        pose = action[:16].reshape(4,4)
+        grip_action = action[16]
+        
+        if self.action_space == "delta":
+            goal_pose = self.get_camera_pose() @ pose
+        else:
+            goal_pose = pose
+        
+        self.move(goal_pose)
+        updated = self.set_grasp(grip_action)
+        if updated:
+            step_fn(self.grasping_steps)
+        step_fn(self.control_steps)

egogym/embodiments/grippers/rum.py

@@ -0,0 +1,6 @@
+from egogym.embodiments.grippers.floating_gripper import FloatingGripper
+
+class RUMGripper(FloatingGripper):
+    def __init__(self, control_steps=200, grasping_steps=2000, action_space="delta", np_random=None):
+        super().__init__(control_steps, grasping_steps, action_space, np_random=np_random)
+        self.name = "rum"

egogym/embodiments/robot.py

@@ -0,0 +1,95 @@
+import mujoco
+from functools import wraps
+import numpy as np
+from egogym.utils import get_pose
+
+
+def requires_mj(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if not hasattr(self, 'model') or not hasattr(self, 'data') or self.model is None or self.data is None:
+            raise RuntimeError(f"Must call setup_mj before using the {self.__class__.__name__}.")
+        return func(self, *args, **kwargs)
+    return wrapper
+
+
+class Robot():    
+    def __init__(self, control_steps=50, grasping_steps=1, action_space="delta", np_random=None):
+        self.name = None
+        self.model = None
+        self.data = None
+        self.control_steps = control_steps
+        self.grasping_steps = grasping_steps
+        self.action_space = action_space
+        self.np_random = np_random
+        self.camera_names = ["egocentric", "exocentric"]
+        self.rgb_renderers = {}
+        self.add_noise_init = True
+
+    def additional_obs_spaces(self):
+        return {}
+    
+    def get_additional_obs(self):
+        return {}
+    
+    def action_space_shape(self):
+        return (0,)
+    
+    def setup_mj(self, model, data):
+        self.model = model
+        self.data = data
+
+    @requires_mj
+    def get_camera_pose(self):
+        return get_pose(self.data, "egocentric", "camera")
+    
+    
+    @requires_mj
+    def setup_renderers(self, render_width=960, render_height=720):
+        if len(self.rgb_renderers) > 0:
+            del self.rgb_renderers
+        self.rgb_renderers = {}
+        for camera_name in self.camera_names:
+            self.rgb_renderers[camera_name] = mujoco.Renderer(self.model, height=render_height, width=render_width)
+    
+    @requires_mj
+    def get_camera_view(self, camera_name):
+        scene_opt = mujoco.MjvOption()
+        self.rgb_renderers[camera_name].update_scene(self.data, camera_name, scene_opt)
+        rgb = self.rgb_renderers[camera_name].render()
+        return rgb
+    
+    @requires_mj
+    def get_tcp_pose(self):
+        return get_pose(self.data, "grasping_center", "site")
+    
+    @requires_mj
+    def get_grasp(self):
+        return np.array(self.data.qpos[self.model.actuator("fingers_actuator").id])
+    
+    @requires_mj
+    def get_grasped_bodies(self):
+        contact_points = self.data.contact
+        body_names = set()
+        for contact in contact_points:
+            body1_id = self.model.geom_bodyid[contact.geom1]
+            body1 = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_BODY, body1_id)
+            body2_id = self.model.geom_bodyid[contact.geom2]
+            body2 = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_BODY, body2_id)
+            if body1 is None or body2 is None:
+                continue
+            if "left" in body1 or "right" in body1:
+                body_names.add(body2)
+            if "left" in body2 or "right" in body2:
+                body_names.add(body1)
+        return body_names
+    
+
+    @requires_mj
+    def set_grasp(self, grip_action):
+        desired_grasp = -255.0 * grip_action
+        self.data.ctrl[self.model.actuator("fingers_actuator").id] = desired_grasp.item()
+        if desired_grasp != self.last_grasp:
+            self.last_grasp = desired_grasp
+            return 1
+        return 0