egogym 0.1.0__py3-none-any.whl

egogym/wrappers/unprivileged_molmo.py

@@ -0,0 +1,88 @@
+import asyncio
+import numpy as np
+import threading
+from egogym.utils import pixel_to_world
+from egogym.misc.molmo_client import MolmoClient
+
+class UnprivilegedMolmo:
+    
+    def __init__(self, env, host="localhost", port=8765):
+        self.env = env
+        self.unprivileged_T_world_object = None
+        self.molmo_client = MolmoClient(host=host, port=port)
+        self._loop = None
+        self._loop_thread = None
+        self._connected = False
+        
+    def _start_event_loop(self):
+        if self._loop is None:
+            self._loop = asyncio.new_event_loop()
+            self._loop_thread = threading.Thread(target=self._run_event_loop, daemon=True)
+            self._loop_thread.start()
+    
+    def _run_event_loop(self):
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+        
+    def _ensure_connected(self):
+        if not self._connected:
+            try:
+                self._start_event_loop()
+                future = asyncio.run_coroutine_threadsafe(self.molmo_client.connect(), self._loop)
+                future.result(timeout=10)
+                self._connected = True
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to connect to Molmo server at {self.molmo_client.uri}. "
+                    f"Make sure the Molmo server is running with: "
+                    f"'python baselines/rum/molmo_server.py'. "
+                    f"Original error: {e}"
+                ) from e
+    
+    def _infer_point_sync(self, rgb: np.ndarray, object_name: str) -> np.ndarray:
+        self._ensure_connected()
+        future = asyncio.run_coroutine_threadsafe(
+            self.molmo_client.infer_point(rgb, object_name=object_name), 
+            self._loop
+        )
+        return future.result(timeout=30)
+        
+    def reset(self, **kwargs):
+        obs, info = self.env.reset(**kwargs)
+        
+        robot = self.env.unwrapped.robot
+        robot.rgb_renderers[robot.camera_names[0]].enable_depth_rendering()
+        robot.rgb_renderers[robot.camera_names[0]].update_scene(robot.data, robot.camera_names[0])
+        depth = robot.rgb_renderers[robot.camera_names[0]].render()
+        robot.rgb_renderers[robot.camera_names[0]].disable_depth_rendering()
+
+        point_norm = self._infer_point_sync(obs["rgb_ego"], obs["object_name"])
+        x_norm, y_norm = point_norm
+        x = int(x_norm * self.env.unwrapped.render_width)
+        y = int(y_norm * self.env.unwrapped.render_height)
+        depth_value = depth[y, x] + 0.03
+        self.unprivileged_T_world_object = pixel_to_world(x, y, depth_value, self.env.unwrapped.model, self.env.unwrapped.data, robot.camera_names[0], self.env.unwrapped.render_width, self.env.unwrapped.render_height)
+
+        return obs, info
+    
+    def step(self, action):
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        
+        object_pose = obs["object_pose"].reshape(4,4)
+        object_pose[:3, 3] = self.unprivileged_T_world_object
+        obs["object_pose"] = object_pose.flatten()
+
+        return obs, reward, terminated, truncated, info
+    
+    def close(self):
+        if self._connected:
+            future = asyncio.run_coroutine_threadsafe(self.molmo_client.disconnect(), self._loop)
+            future.result(timeout=5)
+            self._connected = False
+        if self._loop is not None:
+            self._loop.call_soon_threadsafe(self._loop.stop)
+            self._loop_thread.join(timeout=2)
+        self.env.close()
+    
+    def __getattr__(self, name):
+        return getattr(self.env, name)

egogym/wrappers/unprivileged_moondream.py

@@ -0,0 +1,121 @@
+import os
+import numpy as np
+import time
+import io
+import base64
+import requests
+from PIL import Image
+from egogym.utils import pixel_to_world
+
+MOONDREAM_API_URL = "https://api.moondream.ai/v1/point"
+
+
+class UnprivilegedMoondream:
+    
+    def __init__(self, env, api_key=None):
+        self.env = env
+        self.unprivileged_T_world_object = None
+        
+        if api_key is None:
+            api_key = os.environ.get("MOONDREAM_API_KEY")
+        
+        self.api_key = api_key
+    
+    def _call_api_with_retry(self, image_base64: str, object_name: str, max_retries=3, base_delay=2):
+        """Call Moondream API with retry logic for transient errors."""
+        headers = {
+            "Content-Type": "application/json",
+            "X-Moondream-Auth": self.api_key
+        }
+        
+        payload = {
+            "image_url": f"data:image/jpeg;base64,{image_base64}",
+            "object": object_name
+        }
+        
+        for attempt in range(max_retries):
+            try:
+                response = requests.post(MOONDREAM_API_URL, headers=headers, json=payload, timeout=30)
+                
+                if response.status_code == 200:
+                    return response.json()
+                elif response.status_code in [503, 429] and attempt < max_retries - 1:
+                    delay = base_delay * (2 ** attempt)
+                    print(f"API error {response.status_code}: {response.text}. Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})")
+                    time.sleep(delay)
+                    continue
+                else:
+                    response.raise_for_status()
+            except requests.exceptions.RequestException as e:
+                if attempt < max_retries - 1:
+                    delay = base_delay * (2 ** attempt)
+                    print(f"Request error: {e}. Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})")
+                    time.sleep(delay)
+                    continue
+                else:
+                    raise
+        
+        return None
+    
+    def _infer_point_sync(self, rgb: np.ndarray, object_name: str) -> np.ndarray:
+        """Use Moondream API to locate object in image and return normalized coordinates."""
+        if rgb.dtype == np.float32 or rgb.dtype == np.float64:
+            rgb = (rgb * 255).astype(np.uint8)
+        image = Image.fromarray(rgb)
+        
+        try:
+            buf = io.BytesIO()
+            image.save(buf, format='JPEG')
+            image_bytes = buf.getvalue()
+            image_base64 = base64.b64encode(image_bytes).decode('utf-8')
+            
+            response = self._call_api_with_retry(image_base64, object_name)
+            
+            if response and "points" in response and len(response["points"]) > 0:
+                point = response["points"][0]
+                x_norm = float(point["x"])
+                y_norm = float(point["y"])
+                
+                x_norm = max(0.0, min(1.0, x_norm))
+                y_norm = max(0.0, min(1.0, y_norm))
+                
+                return np.array([x_norm, y_norm], dtype=np.float32)
+            
+            return np.array([0.5, 0.5], dtype=np.float32)
+            
+        except Exception as e:
+            print(f"Moondream API error: {e}")
+            return np.array([0.5, 0.5], dtype=np.float32)
+
+    def reset(self, **kwargs):
+        obs, info = self.env.reset(**kwargs)
+        
+        robot = self.env.unwrapped.robot
+        robot.rgb_renderers[robot.camera_names[0]].enable_depth_rendering()
+        robot.rgb_renderers[robot.camera_names[0]].update_scene(robot.data, robot.camera_names[0])
+        depth = robot.rgb_renderers[robot.camera_names[0]].render()
+        robot.rgb_renderers[robot.camera_names[0]].disable_depth_rendering()
+
+        point_norm = self._infer_point_sync(obs["rgb_ego"], obs["object_name"])
+        x_norm, y_norm = point_norm
+        x = int(x_norm * self.env.unwrapped.render_width)
+        y = int(y_norm * self.env.unwrapped.render_height)
+        depth_value = depth[y, x] + 0.03
+        self.unprivileged_T_world_object = pixel_to_world(x, y, depth_value, self.env.unwrapped.model, self.env.unwrapped.data, robot.camera_names[0], self.env.unwrapped.render_width, self.env.unwrapped.render_height)
+
+        return obs, info
+    
+    def step(self, action):
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        
+        object_pose = obs["object_pose"].reshape(4,4)
+        object_pose[:3, 3] = self.unprivileged_T_world_object
+        obs["object_pose"] = object_pose.flatten()
+
+        return obs, reward, terminated, truncated, info
+    
+    def close(self):
+        self.env.close()
+    
+    def __getattr__(self, name):
+        return getattr(self.env, name)

egogym-0.1.0.dist-info/METADATA

@@ -0,0 +1,52 @@
+Metadata-Version: 2.4
+Name: egogym
+Version: 0.1.0
+Summary: EgoGym: A robotics environment for egocentric control tasks
+License: MIT
+Project-URL: Homepage, https://github.com/omarrayyann/EgoGym
+Project-URL: Repository, https://github.com/omarrayyann/EgoGym
+Keywords: robotics,reinforcement-learning,mujoco,gymnasium
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: mujoco
+Requires-Dist: gymnasium
+Requires-Dist: opencv-python
+Requires-Dist: scipy
+Requires-Dist: gdown
+Requires-Dist: numpy
+Requires-Dist: torch
+Requires-Dist: transformers
+Requires-Dist: tqdm
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: flake8; extra == "dev"
+
+# EgoGym
+
+EgoGym is a lightweight benchmark suite for egocentric robot policies
+
+```python
+env = gym.make(
+    "Egogym-Pick-v0",           # Options: "Egogym-Pick-v0", "Egogym-Open-v0", "Egogym-Close-v0"
+    robot="rum",                # Options: "rum", "droid"
+    action_space="delta",       # Options: "delta", "absolute"
+    num_objs=5,                 # Options: 1-5
+)
+```
+
+**Actions (17-dim):**
+- `[0-15]`: flattened 4x4 transformation
+- `[16]`: continuous gripper (0=open, 1=close)
+
+**Coordinate Frame:**
+- **x**: Right of camera
+- **y**: Up of camera
+- **z**: Backward of camera

egogym-0.1.0.dist-info/RECORD

@@ -0,0 +1,83 @@
+baselines/pi_policy.py,sha256=mT-rJJ--Kp0WxVdWtPHgB0Zd_CS-2U8Cilk4rrf3fFY,4368
+baselines/rum_policy.py,sha256=YTw-3O6sLAy5UZ9begWlY3VwWI-PMIlxIMWF8_Y1pxc,4549
+baselines/rum/__init__.py,sha256=6SyMXPP_WUscCLtrgntS-Oq8FDivQ5j9UJHiX6rYKb4,27
+baselines/rum/policy.py,sha256=KBIT19nzhhq6PB-pxzcxSCOZ4jUN2hfxiuT58nItz_0,10775
+baselines/rum/loss_fns/__init__.py,sha256=arlX1c2M3S-2n0BHRhhYCu7tn3i0qYZchTeMtKsVoh8,1163
+baselines/rum/loss_fns/abstract_loss_fn.py,sha256=CmCAGmgJWKpdPv2e8O6uDAq3a5ktIs5MukMQmOu8bNc,291
+baselines/rum/loss_fns/diffusion_policy_loss_fn.py,sha256=Ef-LDQh3EvFG50U4C1n-Zpk3MCqQ1W2UtVurTX9uM3g,3863
+baselines/rum/loss_fns/rvq_loss_fn.py,sha256=Li4j4Eb-wy5qF0Sjk5f3EjOWtVObkoXuaFbxusXcCc4,3424
+baselines/rum/loss_fns/vqbet_loss_fn.py,sha256=ZN4mPdSQ9Ngx4q9jjBfZa5f8myhrjhASpBAOGplpHWY,7326
+baselines/rum/models/__init__.py,sha256=Il5Q9ATdX8yXqVxtP_nYqUhExzxPC_qk_WXQ_4h0exg,16
+baselines/rum/models/custom.py,sha256=9oTuQI0udxRnotnQQqRDMZkMFN6VFSsF2Y8KMtf1_cw,1073
+baselines/rum/models/bet/__init__.py,sha256=Snvarck4ui4gv05Zx1ovSfTadkKL4s3sFjpIHj840ZA,186
+baselines/rum/models/bet/bet.py,sha256=dGuyUo9gEt2rM-rDPp_yYmrMUC333Hydq2jSeMdqqsk,13660
+baselines/rum/models/bet/gpt.py,sha256=rzJK9ilz0U3Ca2d68A0mHklFO1UaDtzZu3XX6ez_z50,10739
+baselines/rum/models/bet/tokenized_bet.py,sha256=kL7YF-rqkHj9Dv_cND1ObRydWbtp19-QpMnNifw2V8s,19168
+baselines/rum/models/bet/utils.py,sha256=OTJTEssmCPvhAa6jbCYar-ol2XS3c0xSBsvPsJzab8k,5665
+baselines/rum/models/bet/vqbet.py,sha256=zj2LWwKkDoIk7SCIN9WT3BfPrMWfOG2IIhVNLzAuwr4,17021
+baselines/rum/models/bet/vqvae/__init__.py,sha256=70gnPjHY1za7k-nViX9GZBrVG87eLN4FKIEcIKKTpIE,203
+baselines/rum/models/bet/vqvae/residual_vq.py,sha256=BgMyTkWtXbCOTJ2tREwTgjugxcLk6zzX6qHU9fgg7bY,10842
+baselines/rum/models/bet/vqvae/vector_quantize_pytorch.py,sha256=d9tSdCjlYp3qtqXr7Zya6OnnTbGjtOApeu4ksVdPd0s,38223
+baselines/rum/models/bet/vqvae/vqvae.py,sha256=QfAOs8fBftuT1CwDPOI_nl2SKggW6Q_EPto0FM8M5Rs,11556
+baselines/rum/models/bet/vqvae/vqvae_utils.py,sha256=pUpgXb7Jt_H9Ou7rrtqdfj_tYD10oZScSur2JtH2ioc,968
+baselines/rum/models/encoders/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+baselines/rum/models/encoders/abstract_base_encoder.py,sha256=TpKMrqQBwgwS9w5X7LkHK7iunAlUZvBM_ERJTnbwOQE,1813
+baselines/rum/models/encoders/identity.py,sha256=y7_tVFKfRUl0CIydvisEvwe3_30IWKGItALf3g4VPd0,1136
+baselines/rum/models/encoders/timm_encoders.py,sha256=LvKVOwq1lD0Q-b7m2FoVXKePz-xN23Fmg0F3bwCx8A4,2519
+baselines/rum/models/policies/diffusion_policy.py,sha256=3s5tyUdd1vOjCYcbmFJpzrforTL0P3x5UlaHYWJgSIc,32096
+baselines/rum/models/policies/open_loop.py,sha256=hirJ03CpktSqKB2GcGrM-5eA7TpgBqK1VIZ3l55FSrI,4175
+baselines/rum/models/policies/simple_open_loop.py,sha256=3HvkaC2uji0gvI1HswuNidVqOwd0G0Dls73ZiZKiDB4,3641
+baselines/rum/molmo/server.py,sha256=cR4TGSYslwNTfnU6SF-ixc7X5i-RqHDzFGRF5b42BTU,4511
+baselines/rum/utils/__init__.py,sha256=x3Yhx4jqkPQtdcVyBuvBA_8BgWWb_-IKhKIGeLkYBv4,7094
+baselines/rum/utils/action_transforms.py,sha256=kSi6UPn8fM1NQSviHvSeXNljNOuVZECwaFdgyp-Yx4o,674
+baselines/rum/utils/decord_transforms.py,sha256=b6o3b13ilwfrsxAxVoHgDbABk4lNKLLw0RXiR4l89Oc,4680
+baselines/rum/utils/rpc.py,sha256=46hJFS79kUvyX5-rhm5hYLz1mBS7hF_sC5UkaaF_dJQ,8278
+baselines/rum/utils/schedulers.py,sha256=Rjl-HJdpehkBy_zoJLy5u0YDdIweO-7FUFSMQgSVZ3o,2508
+baselines/rum/utils/trajectory_vis.py,sha256=nUXNV8cwra4k91C3loZd_Lg4l5leG1GOkSGRqGCiNAU,4310
+baselines/rum/utils/zmq_utils.py,sha256=ZSsNaDqpZsGFZRWamRrT0SwuiPwru007kOUcHO-oLjg,9719
+egogym/__init__.py,sha256=CuZXrmju4K0sGveackQCO5ySeiLv23jqVY0Nx25eqes,295
+egogym/egogym.py,sha256=FLhqcIa-H40zmkS85ODuqlixDegXMnVqSbjloVY5JCk,4693
+egogym/evaluate.py,sha256=5dz24UXmfQDBU_Yn7bftRzNrCIhgbLCvQJDm3DqwNzA,7971
+egogym/utils.py,sha256=Yv_bnTIa1wKh_o--u_JwczLqrvmJQnYHm-21N3GBGyw,40723
+egogym/assets/constants.py,sha256=L9Pf5TvmHG4OUSiy3KTb-nWC0yI57T0aMrgsiAZbjqM,29200
+egogym/components/__init__.py,sha256=c8B_uNbJ7iKUFTRoDjJgveapvnyS8QYJtD5VVKVQRwc,45
+egogym/components/object.py,sha256=gJXg6_siNJig7MReTxjZM8soh17d9EhcyZlv4wEkbZM,3655
+egogym/embodiments/__init__.py,sha256=-apnOCtVcTlgH0WzOKRlWG3-4UJkkJcjZaG-ybzKUQ4,265
+egogym/embodiments/robot.py,sha256=FZ4--54DhvvOs7GEKdd-nUiTYz2x_6lh3w3oMD3jIkg,3238
+egogym/embodiments/arms/__init__.py,sha256=GlfQkjW9nuIuvUbprUi0MzDEqc6knh6HdY4EjyMfQK8,80
+egogym/embodiments/arms/arm.py,sha256=SBk9oJwhCaLNqA9OeNfi88905_AHasTa8M5eb0ZX5lc,2252
+egogym/embodiments/arms/droid.py,sha256=piZLBt7GBQwHv3alHdRvWAkDiuf6kVbnsCpwfi7srzw,2120
+egogym/embodiments/grippers/__init__.py,sha256=uQlGAp2wPStYurIvHKmBx85g2r9XABpzcSRC6qoA2FA,119
+egogym/embodiments/grippers/floating_gripper.py,sha256=NdaBs_PEoXzbKhJ1EI-252x_YrbtziTrJxS-iqo9Pi8,2219
+egogym/embodiments/grippers/rum.py,sha256=gc7RXQh_jXxTPiPEP06W1b21nhZvwE0aMy48binGSLA,327
+egogym/managers/__init__.py,sha256=j3m90-JRo5322fBPy4VHx9LNG_Xivl9upYolIvYJhcI,91
+egogym/managers/objects_managers.py,sha256=HIf2BRoSBR5oKQVo-NT7uNoevSTUBFXHJy010GZqgbk,1155
+egogym/managers/textures_manager.py,sha256=NsRL0aPHpky5aDiZazyP1bFozebBn-Thcr_pFaBMt4s,654
+egogym/misc/molmo_client.py,sha256=Fz_Ix8CjxhIVbzTbyPwZmFS1I0QVZIjrwp5t8zypPvE,1408
+egogym/misc/molmo_server.py,sha256=cQ3EJSrH4ptfUHWNQwJN5vAHQiUnEFo7fVh_LdYOLVQ,6171
+egogym/policies/__init__.py,sha256=X1q8mB870ZcNcuhjAkdXOT4bK5P5N-4VsFd6car0m3I,36
+egogym/policies/base_policy.py,sha256=nkLchlGKPIyHlT9aEuaSm5HynPRyu0xV2-mVG0t_WtU,228
+egogym/scripts/analayze.py,sha256=SVWRDZvm1cGK4JUpXXDNpLl52Q9YqDIVNUYk4h23MdE,30109
+egogym/scripts/plot.py,sha256=RkMhA0ULcLnwTCawZPfJ9KlwyB1Eb8zzFmxiFxOAImk,2786
+egogym/scripts/plot_correlation.py,sha256=hyvh2ruicU8wp6Jupuok3bxyVKfjFhV8akdNwtcL4Kw,13878
+egogym/scripts/plot_correlation_hardcoded.py,sha256=ZP93MT-kP2G2yv15-8SLySqz0OkQPxWllNRx51K2ZkE,11325
+egogym/scripts/plot_failure.py,sha256=CdKxWEmghhKoHEmprQ9-YBgMeD6P1-shSgbfVMvi7GU,9079
+egogym/scripts/plot_failure_hardcoded.py,sha256=ozFk6ETgcKmjvJwdKpoSfC_Nj0-NJn94a9gQ1tzJMh8,6551
+egogym/scripts/plot_failure_vlm.py,sha256=imY6yLQ2vqgDMg0kM1LnsnVUWbL_3rOGlBKAdRya6Xg,8916
+egogym/scripts/plot_failure_vlm_hardcoded.py,sha256=BlHIgsf6SMmCKuTF0XMssr9MaPVsdazb-eyF_yS0vMs,5432
+egogym/scripts/plot_line.py,sha256=fjT9JzE7meTy2YRBkzbCTV3tslXm4UAZMi22r4ZFSmc,9831
+egogym/scripts/plot_line_hardcoded.py,sha256=hDsfEGc-_CzONYy_VcWHqUNG4BqI_94-D6fIakfQfdw,8779
+egogym/scripts/plot_pi0_bars.py,sha256=TRSF95ggAwtPCxiSmRGm_xCcF_A2BDIJMNEP6btJfVA,4856
+egogym/tasks/close.py,sha256=snrd-FngKbgXoWijuaEX8eW_KtliefhTucez8xhqi-E,3739
+egogym/tasks/open.py,sha256=MfYJH4RJYhD8xFqbC4TmalorCCwfoqO2LnxZbwwx0L8,3761
+egogym/tasks/pick.py,sha256=ecdYB-c9-JLEsAYK_7655VnIdFOC2yP2HsDev04ufr0,5737
+egogym/wrappers/__init__.py,sha256=7qlwlbL9KQjGHiRUb7GBr7zz5W9o2b6HedsPNmUeh58,762
+egogym/wrappers/episode_monitor.py,sha256=jI-O7HSOo7sJZGaaUXe9ZJ4q0Pyi13FcHUpuu430sLk,10923
+egogym/wrappers/unprivileged_chatgpt.py,sha256=inkpa9gZvqbDiCVhWRDUHnv1TVzmED0Sv6jgGfWYWuE,6356
+egogym/wrappers/unprivileged_gemini.py,sha256=DGnSZ9gZVTyrH5iFRcyKtO-yXm6KaeW0vFfjzYDuumM,6252
+egogym/wrappers/unprivileged_molmo.py,sha256=YCTa_O3hv17tX_WxQ1WQ8RIAIIX4gRrR6UTTfIdiOPM,3587
+egogym/wrappers/unprivileged_moondream.py,sha256=iRNYYURcDFQkCYJEnr6fjPQUgaCAtvmS2ol4J3Rcytk,4762
+egogym-0.1.0.dist-info/METADATA,sha256=kP6g0e3I5tguQyKtL0YLv50qGCoaO7lTnnYtlH6MADE,1651
+egogym-0.1.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+egogym-0.1.0.dist-info/top_level.txt,sha256=JLnBIwPgLWhzMxIhuhpBvdPnG-nWQ7QURdn5Qgml6lw,17
+egogym-0.1.0.dist-info/RECORD,,

egogym-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.10.2)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

egogym-0.1.0.dist-info/top_level.txt

@@ -0,0 +1,2 @@
+baselines
+egogym