egogym 0.1.0__py3-none-any.whl

egogym/scripts/plot_pi0_bars.py

@@ -0,0 +1,169 @@
+import os
+import numpy as np
+import pandas as pd
+import altair as alt
+from scipy.stats import beta
+
+# Register custom font for PNG export
+alt.themes.register('custom_theme', lambda: {
+    'config': {
+        'title': {'font': 'Produkt'},
+        'axis': {'labelFont': 'Produkt', 'titleFont': 'Produkt'},
+        'legend': {'labelFont': 'Produkt', 'titleFont': 'Produkt'},
+        'mark': {'font': 'Produkt'},
+        'text': {'font': 'Produkt'},
+    }
+})
+alt.themes.enable('custom_theme')
+
+BASE_DIR = "logs"
+REWARD_THRESHOLD = 0.03
+
+def compute_success_from_csv(csv_path):
+    df = pd.read_csv(csv_path, sep="\t")
+    successes = (df["max_reward"] > REWARD_THRESHOLD).sum()
+    total = len(df)
+    return successes, total
+
+
+def plot_pi0_bars():
+    model_folder = "pi0"
+    num_objects = [1, 2, 3, 4, 5]
+    
+    rows = []
+    
+    for n_obj in num_objects:
+        possible_folders = [
+            f"{n_obj}_objects",
+            f"{n_obj}_object",
+            f"{n_obj}-objects",
+            f"{n_obj}-object",
+        ]
+        
+        csv_path = None
+        for folder in possible_folders:
+            # First try direct path
+            candidate = os.path.join(BASE_DIR, model_folder, folder, "log.csv")
+            if os.path.exists(candidate):
+                csv_path = candidate
+                break
+            
+            # Try nested evaluation folder structure
+            folder_path = os.path.join(BASE_DIR, model_folder, folder)
+            if os.path.exists(folder_path) and os.path.isdir(folder_path):
+                for subdir in os.listdir(folder_path):
+                    subdir_path = os.path.join(folder_path, subdir)
+                    if os.path.isdir(subdir_path):
+                        candidate = os.path.join(subdir_path, "log.csv")
+                        if os.path.exists(candidate):
+                            csv_path = candidate
+                            break
+                if csv_path:
+                    break
+        
+        if csv_path is None:
+            print(f"Missing data: π-0.5, {n_obj} objects")
+            continue
+        
+        s, t = compute_success_from_csv(csv_path)
+        
+        # Beta posterior
+        a, b = 1 + s, 1 + (t - s)
+        mean = 100 * a / (a + b)
+        lo = 100 * beta.ppf(0.025, a, b)
+        hi = 100 * beta.ppf(0.975, a, b)
+        
+        print(f"π-0.5 | {n_obj} objects: {s}/{t} = {mean:.1f}%")
+        
+        rows.append({
+            "num_objects": n_obj,
+            "mean": mean,
+            "lo": lo,
+            "hi": hi,
+            "successes": s,
+            "total": t
+        })
+    
+    df = pd.DataFrame(rows)
+    
+    # Create bar chart
+    bars = alt.Chart(df).mark_bar(
+        color='#E0BE16',
+        width=60
+    ).encode(
+        x=alt.X(
+            'num_objects:O',
+            title='Number of Objects',
+            axis=alt.Axis(
+                labelFontSize=20,
+                titleFontSize=24,
+                labelAngle=0,
+                titlePadding=15
+            )
+        ),
+        y=alt.Y(
+            'mean:Q',
+            title='Success Rate (%)',
+            scale=alt.Scale(domain=[0, 100]),
+            axis=alt.Axis(
+                labelFontSize=20,
+                titleFontSize=24,
+                titlePadding=14,
+                grid=True,
+                gridOpacity=0.3
+            )
+        ),
+        tooltip=[
+            alt.Tooltip('num_objects:O', title='Objects'),
+            alt.Tooltip('mean:Q', title='Success Rate (%)', format='.1f'),
+            alt.Tooltip('successes:Q', title='Successes'),
+            alt.Tooltip('total:Q', title='Total')
+        ]
+    )
+    
+    # Create error bars for 95% CI
+    error_bars = alt.Chart(df).mark_errorbar(
+        ticks=True,
+        thickness=2
+    ).encode(
+        x=alt.X('num_objects:O'),
+        y=alt.Y('lo:Q', title=''),
+        y2=alt.Y2('hi:Q')
+    )
+    
+    # Add text labels on top of bars
+    text = alt.Chart(df).mark_text(
+        dy=-10,
+        fontSize=16,
+        fontWeight='bold'
+    ).encode(
+        x=alt.X('num_objects:O'),
+        y=alt.Y('mean:Q'),
+        text=alt.Text('mean:Q', format='.1f')
+    )
+    
+    # Combine layers
+    chart = (bars + error_bars + text).properties(
+        width=500,
+        height=400,
+        title={
+            'text': '           π-0.5 Success Rate by Number of Objects',
+            'fontSize': 24,
+            'anchor': 'start',
+            'dx': 40,
+            'dy': -10
+        },
+        padding={'left': 10, 'right': 10, 'top': 40, 'bottom': 40}
+    ).configure_view(
+        strokeWidth=0
+    )
+    
+    return chart
+
+
+if __name__ == "__main__":
+    chart = plot_pi0_bars()
+    chart.save("pi0_bars.html")
+    chart.save("pi0_bars.png", scale_factor=3)
+    chart.save("pi0_bars.pdf", scale_factor=3)
+    print("\nPlot saved to: pi0_bars.html, pi0_bars.png, and pi0_bars.pdf")

egogym/tasks/close.py

@@ -0,0 +1,84 @@
+import os
+import cv2
+import mujoco
+from gymnasium.spaces import Box
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
+from egogym.egogym import Egogym
+from egogym.utils import include_in_scene, position_sampler, make_objects_manager
+import egogym.assets.constants as constants
+
+class CloseTask(Egogym):
+
+    def __init__(self, robot="rum", action_space="delta", render_mode=None, render_size=(960,720), num_objs=1, seed=None, objects_set=None):
+        super().__init__(robot=robot, action_space=action_space, render_mode=render_mode, render_size=render_size, seed=seed)
+        self.num_objs = num_objs
+        if objects_set is not None:
+            self.objects_manager = make_objects_manager(objects_set, self.np_random, shuffle=False)
+        else:
+            self.objects_manager = make_objects_manager(constants.all_close_objects_set, self.np_random)
+        self.observation_space["handle_pose"] = Box(low=-np.inf, high=np.inf, shape=(16,), dtype=np.float32)
+
+    def make_task_scene(self):
+        with open(f"{os.path.dirname(__file__)}/../assets/scenes/open.xml", "r") as f:
+            scene_xml = f.read()
+
+        self.object = self.objects_manager.sample()
+        scene_xml = self.object.add_to_scene_xml(scene_xml)
+        scene_xml = include_in_scene(scene_xml, f"{os.path.dirname(__file__)}/../assets/embodiments/{self.robot.name}/model_open.xml")
+        return scene_xml
+
+    def get_obs(self):
+        obs = self.get_base_obs()
+        obs["handle_pose"] = self.object.get_handle_pose(self.data).reshape(16).astype(np.float32)
+        return obs
+
+    def compute_reward(self) -> float:
+        perecentage_opened = self.object.get_perecenttage_opened(self.model, self.data)
+        return max(0.1-perecentage_opened, 0.0)
+
+    def reset(self, seed=None, options=None):
+        super().reset(seed=seed, options=options)
+        scene_xml_string = self.make_task_scene()
+        self.grasped_bodies = set()
+        self.grasping_object = False
+        
+        gripper_init_pose = np.eye(4)
+        gripper_init_pose[:3, :3] = R.from_euler("xyz", np.array([1.3, 0.000, 0.000])).as_matrix()
+        self.setup_mujoco(scene_xml_string)
+        self.env_step()
+        self.object.open(self.model, self.data)
+        self.env_step()
+        handle_pose = self.object.get_handle_pose(self.data)
+        camera_pos = handle_pose[0:3, 3].copy() + np.array([-0.1, -0.8, 0.35])
+        gripper_init_pose[:3, 3] = handle_pose[0:3, 3] + np.array([0.00, -0.6, 0.2])
+        self.model.cam_pos[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_CAMERA, "exocentric")] = camera_pos
+        self.robot.prepare(gripper_init_pose, self.env_step)
+        self.env_step(10)
+        self.initial_robot_pose = self.robot.get_camera_pose()
+
+        self.enable_sleeping_islands()
+
+        observation = self.get_obs()
+        info = {"object_name": self.object.name}
+
+        return observation, info
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = super().step(action)
+        self.grasped_bodies.update(self.robot.get_grasped_bodies())
+        grasped_bodies_list = list(self.grasped_bodies)
+        if f"{self.object.name}_object" in grasped_bodies_list:
+            self.grasping_object = True
+        info = {
+            "grasped_bodies": grasped_bodies_list,
+            "object_name": self.object.name,
+            "initial_robot_pose": self.initial_robot_pose,
+            "initial_object_pose": self.object.last_set_pose,
+            "is_grasping": self.robot.get_grasp(),
+            "gripper_current_position": self.robot.get_tcp_pose(),
+            "grasping_object": self.grasping_object,
+            **info
+        }
+        return obs, reward, terminated, truncated, info

egogym/tasks/open.py

@@ -0,0 +1,85 @@
+import os
+import cv2
+import mujoco
+from gymnasium.spaces import Box
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
+from egogym.egogym import Egogym
+from egogym.utils import include_in_scene, position_sampler, make_objects_manager
+import egogym.assets.constants as constants
+
+class OpenTask(Egogym):
+
+    def __init__(self, robot="rum", action_space="delta", render_mode=None, render_size=(960,720), num_objs=1, seed=None, objects_set=None):
+        super().__init__(robot=robot, action_space=action_space, render_mode=render_mode, render_size=render_size, seed=seed)
+        self.num_objs = num_objs
+        if objects_set is not None:
+            self.objects_manager = make_objects_manager(objects_set, self.np_random, shuffle=False)
+        else:
+            self.objects_manager = make_objects_manager(constants.all_open_objects_set, self.np_random)
+        self.observation_space["handle_pose"] = Box(low=-np.inf, high=np.inf, shape=(16,), dtype=np.float32)
+
+    def make_task_scene(self):
+        with open(f"{os.path.dirname(__file__)}/../assets/scenes/open.xml", "r") as f:
+            scene_xml = f.read()
+
+        self.object = self.objects_manager.sample()
+        scene_xml = self.object.add_to_scene_xml(scene_xml)
+        scene_xml = include_in_scene(scene_xml, f"{os.path.dirname(__file__)}/../assets/embodiments/{self.robot.name}/model_open.xml")
+        return scene_xml
+
+    def get_obs(self):
+        obs = self.get_base_obs()
+        obs["handle_pose"] = self.object.get_handle_pose(self.data).reshape(16).astype(np.float32)
+        return obs
+
+    def compute_reward(self) -> float:
+        perecentage_opened = self.object.get_perecenttage_opened(self.model, self.data)
+        return perecentage_opened
+
+    def reset(self, seed=None, options=None):
+        super().reset(seed=seed, options=options)
+        scene_xml_string = self.make_task_scene()
+        self.grasped_bodies = set()
+        self.grasping_object = False
+        
+        gripper_init_pose = np.eye(4)
+        gripper_init_pose[:3, :3] = R.from_euler("xyz", np.array([1.5, 0.000, 0.000])).as_matrix()
+        self.setup_mujoco(scene_xml_string)
+        mujoco.mj_step(self.model, self.data)
+        handle_pose = self.object.get_handle_pose(self.data)
+        camera_pos = handle_pose[0:3, 3].copy()
+        camera_pos += np.array([0.0, -1.0, 0.28])
+        self.model.cam_pos[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_CAMERA, "exocentric")] = camera_pos
+        mujoco.mj_step(self.model, self.data)
+        gripper_init_pose[:3, 3] = handle_pose[:3, 3] + np.array([0.0, -0.6, 0.05])
+        self.robot.prepare(gripper_init_pose, self.env_step)
+        self.env_step()
+        self.initial_robot_pose = self.robot.get_camera_pose()
+
+        self.settle_env()
+        self.enable_sleeping_islands()
+
+        observation = self.get_obs()
+        info = {"object_name": self.object.name}
+
+        return observation, info
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = super().step(action)
+        self.grasped_bodies.update(self.robot.get_grasped_bodies())
+        grasped_bodies_list = list(self.grasped_bodies)
+        if f"{self.object.name}_object" in grasped_bodies_list:
+            self.grasping_object = True
+        info = {
+            "grasped_bodies": grasped_bodies_list,
+            "object_name": self.object.name,
+            "initial_robot_pose": self.initial_robot_pose,
+            "initial_object_pose": self.object.last_set_pose,
+            "is_grasping": self.robot.get_grasp(),
+            "gripper_current_position": self.robot.get_tcp_pose(),
+            "grasping_object": self.grasping_object,
+            **info
+        }
+        return obs, reward, terminated, truncated, info

egogym/tasks/pick.py

@@ -0,0 +1,121 @@
+import os
+from gymnasium.spaces import Box, Text
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
+from egogym.egogym import Egogym
+from egogym.utils import include_in_scene, position_sampler, make_objects_manager, make_textures_manager
+import egogym.assets.constants as constants
+
+class PickTask(Egogym):
+
+    def __init__(self, robot="rum", action_space="delta", render_mode=None, render_size=(960,720), num_objs=1, seed=None, objects_set=None):
+        super().__init__(robot=robot, action_space=action_space, render_mode=render_mode, render_size=render_size, seed=seed)
+        self.spread = 0.22
+        self.num_objs = num_objs
+        if objects_set is not None:
+            self.objects_manager = make_objects_manager(objects_set, self.np_random, shuffle=True)
+        else:
+            self.objects_manager = make_objects_manager(constants.lite_pick_objects_set, self.np_random)
+        self.textures_manager = make_textures_manager([f"wood/{i}.png" for i in range(10)], self.np_random)
+        self.observation_space["object_pose"] = Box(low=-np.inf, high=np.inf, shape=(16,), dtype=np.float32)
+        self.observation_space["object_name"] = Text(max_length=256)
+
+
+    def make_task_scene(self):
+        assets_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets"))
+        
+        with open(os.path.join(assets_dir, "scenes", "pick.xml"), "r") as f:
+            scene_xml = f.read()
+
+        self.object = self.objects_manager.sample()
+        self.secondary_objects = []
+        used_objects = {self.object.name}
+        scene_xml = self.object.add_to_scene_xml(scene_xml)
+        while len(used_objects) < self.num_objs:
+            object = self.objects_manager.sample(random=True)
+            if object.name not in used_objects and object.name.split("_")[0] != self.object.name.split("_")[0]:
+                self.secondary_objects.append(object)
+                used_objects.add(object.name)
+                scene_xml = object.add_to_scene_xml(scene_xml)
+        self.texture = self.textures_manager.sample(random=True)
+        scene_xml = include_in_scene(scene_xml, os.path.join(assets_dir, "embodiments", self.robot.name, "model_pick.xml"))
+        scene_xml = scene_xml.replace("{ASSETS_DIR}", assets_dir)
+        scene_xml = scene_xml.replace("{TEXTURE_PATH}", self.texture)
+        scene_xml = scene_xml.replace("{TABLE_WIDTH}", str(self.spread+0.24))
+        scene_xml = scene_xml.replace("{TABLE_HEIGHT}", str(self.spread+0.24))
+        return scene_xml
+
+    def get_obs(self):
+        obs = self.get_base_obs()
+        obs["object_pose"] = self.object.get_pose(self.data).reshape(16).astype(np.float32)
+        obs["object_name"] = self.object.name.split("_")[0]
+        return obs
+
+    def compute_reward(self) -> float:
+        object_pose = self.object.get_pose(self.data)
+        lift_distance = object_pose[2][3] - self.object.last_set_pose[2][3]
+        return lift_distance
+
+    def reset(self, seed=None, options=None):
+        super().reset(seed=seed, options=options)
+        scene_xml_string = self.make_task_scene()
+        self.grasped_bodies = set()
+        self.grasping_object = False
+        
+        gripper_init_pose = np.eye(4)
+        gripper_init_pose[:3, 3] = np.array([0.00, -0.68, 1.1])
+        gripper_init_pose[:3, :3] = R.from_euler("xyz", np.array([1.3, 0.000, 0.000])).as_matrix()
+        self.setup_mujoco(scene_xml_string)
+        initial_positions = position_sampler(self.np_random, len(self.secondary_objects) + 1, [self.spread,self.spread], self.spread/2, 100, start_position=gripper_init_pose[:2, 3], thickness=self.spread/2.5)
+        pose = np.eye(4)
+        self.initial_robot_pose = self.robot.get_camera_pose()
+        self.env_step(10)
+
+        for i, obj in enumerate([self.object] + self.secondary_objects):
+            z = obj.get_center_to_bottom_z_distance(self.data) + 0.78
+            pose[:2, 3] = initial_positions[i][:2]
+            pose[2, 3] = z
+            pose[:3, :3] = R.from_euler(
+                "x", self.np_random.uniform(0, 2 * np.pi)
+            ).as_matrix()
+            obj.set_pose(self.model, self.data, pose)
+        
+        self.env_step(100)
+        self.settle_env()
+        self.env_step(100)
+        for obj in [self.object] + self.secondary_objects:
+            obj.last_set_pose = obj.get_pose(self.data).copy()
+        
+        self.robot.prepare(gripper_init_pose, self.env_step)
+        self.env_step(100)        
+        self.enable_sleeping_islands()
+
+        observation = self.get_obs()
+        info = {"object_name": self.object.name}
+
+        return observation, info
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = super().step(action)
+        self.grasped_bodies.update(self.robot.get_grasped_bodies())
+        grasped_bodies_list = list(self.grasped_bodies)
+        if f"{self.object.name}_object" in grasped_bodies_list:
+            self.grasping_object = True
+        info = {
+            "grasped_bodies": grasped_bodies_list,
+            "object_name": self.object.name,
+            "texture_name": self.texture,
+            "initial_robot_pose": self.initial_robot_pose,
+            "initial_object_pose": self.object.last_set_pose,
+            "is_grasping": self.robot.get_grasp(),
+            "gripper_current_position": self.robot.get_tcp_pose(),
+            "grasping_object": self.grasping_object,
+            **info
+        }
+
+        if np.linalg.norm(self.robot.get_camera_pose()[:3,3] - self.object.get_pose(self.data)[:3,3]) > np.linalg.norm(self.initial_robot_pose[:3,3] - self.object.last_set_pose[:3,3])*1.5:
+            truncated = True
+            terminated = True
+
+        return obs, reward, terminated, truncated, info