egogym 0.1.0__py3-none-any.whl

egogym/wrappers/__init__.py

@@ -0,0 +1,20 @@
+from egogym.wrappers.episode_monitor import EpisodeMonitor
+from egogym.wrappers.unprivileged_molmo import UnprivilegedMolmo
+from egogym.wrappers.unprivileged_gemini import UnprivilegedGemini
+from egogym.wrappers.unprivileged_moondream import UnprivilegedMoondream
+from egogym.wrappers.unprivileged_chatgpt import UnprivilegedChatGPT
+
+VLM_WRAPPERS = {
+    "molmo": UnprivilegedMolmo,
+    "gemini": UnprivilegedGemini,
+    "moondream": UnprivilegedMoondream,
+    "chatgpt": UnprivilegedChatGPT,
+}
+
+
+def get_vlm_wrapper(vlm_name: str):
+    if vlm_name not in VLM_WRAPPERS:
+        available = ", ".join(f"'{k}'" for k in VLM_WRAPPERS.keys())
+        raise ValueError(f"Unknown VLM option: '{vlm_name}'. Choose from {available}.")
+    return VLM_WRAPPERS[vlm_name]
+

egogym/wrappers/episode_monitor.py

@@ -0,0 +1,282 @@
+import numpy as np
+import csv
+import os
+import cv2
+
+class EpisodeMonitor:
+    
+    def __init__(self, env, logs_dir=None, columns=None, record=False, render_freq=0, num_envs=1):
+        self.env = env
+        self.logs_dir = logs_dir
+        self.record = record
+        self.mujoco_step_counter = 0
+        self.columns = columns or [
+            "episode", "max_reward", "object_name", "texture_name", 
+            "grasped_bodies", "steps", "initial_robot_pose", 
+            "initial_object_pose", "is_grasping",
+            "gripper_current_position", "grasping_object"
+        ]
+        
+        self.is_vector = num_envs > 1
+        self.render_freq = render_freq
+        
+        if self.is_vector:
+            self.num_envs = num_envs
+            self.global_episode_counter = 0
+            self.episode_idx = np.zeros(self.num_envs, dtype=int)
+            self.max_reward = np.zeros(self.num_envs)
+            self.steps = np.zeros(self.num_envs, dtype=int)
+            self.episode_info = [{} for _ in range(self.num_envs)]
+            if self.record:
+                self.video_writers = [None] * self.num_envs
+                self.video_frames = [[] for _ in range(self.num_envs)]
+        else:
+            self.episode_idx = 0
+            self.max_reward = 0
+            self.steps = 0
+            self.episode_info = {}
+            if self.record:
+                self.video_writer = None
+                self.video_frames = []
+        
+        self.csv_writer = None
+        self.csv_file = None
+
+        os.makedirs(self.logs_dir, exist_ok=True)
+        csv_path = os.path.join(self.logs_dir, "log.csv")
+        self.csv_file = open(csv_path, 'w', newline='')
+        self.csv_writer = csv.DictWriter(self.csv_file, fieldnames=self.columns, delimiter='\t')
+        self.csv_writer.writeheader()
+        
+        if self.record:
+            self.video_dir = os.path.join(self.logs_dir, "videos")
+            os.makedirs(self.video_dir, exist_ok=True)
+        
+        self._callbacks_initialized = False
+        self._callbacks_paused = False
+        self._callbacks_paused = False
+    
+    def _setup_callbacks(self):
+        if self._callbacks_initialized or self.render_freq <= 0:
+            return
+        
+        if self.is_vector:
+            # AsyncVectorEnv doesn't provide direct access to individual envs
+            # Callbacks cannot be set up for vectorized environments
+            pass
+        else:
+            def callback():
+                self._capture_frame_single()
+            actual_env = self.env.unwrapped
+            actual_env._render_callback = callback
+            actual_env._render_freq = self.render_freq
+            actual_env._should_render = actual_env.render_mode == 'human'
+        
+        self._callbacks_initialized = True
+    
+    def _capture_frame_single(self):
+        if self.render_freq <= 0 or self._callbacks_paused:
+            return
+        
+        try:
+            actual_env = self.env.unwrapped
+            
+            if actual_env._should_render:
+                actual_env.render()
+            
+            if self.record:
+                frame_ego = actual_env.robot.get_camera_view(actual_env.robot.camera_names[0])
+                frame_exo = actual_env.robot.get_camera_view(actual_env.robot.camera_names[1])
+                frame = np.concatenate((frame_ego, frame_exo), axis=1)
+                frame = self._process_frame(frame)
+                self.video_frames.append(frame)
+        except Exception as e:
+            pass
+    
+    def _capture_frame_vectorized(self, env_idx):
+        if self.render_freq <= 0 or self._callbacks_paused:
+            return
+        
+        try:
+            env = self.env.envs[env_idx]
+            actual_env = env.unwrapped
+            
+            if actual_env._should_render:
+                actual_env.render()
+            
+            if self.record:
+                frame_ego = actual_env.robot.get_camera_view(actual_env.robot.camera_names[0])
+                frame_exo = actual_env.robot.get_camera_view(actual_env.robot.camera_names[1])
+                frame = np.concatenate((frame_ego, frame_exo), axis=1)
+                frame = self._process_frame(frame)
+                self.video_frames[env_idx].append(frame)
+        except Exception as e:
+            pass
+    
+    def _process_frame(self, frame):
+        if isinstance(frame, np.ndarray):
+            if frame.dtype != np.uint8:
+                frame = (frame * 255).astype(np.uint8) if frame.max() <= 1.0 else frame.astype(np.uint8)
+            if frame.shape[-1] == 3:
+                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+        return frame
+    
+    def reset(self, **kwargs):
+        self._callbacks_paused = True
+        
+        if self.is_vector:
+            options = kwargs.get('options', {})
+            reset_mask = options.get('reset_mask', np.ones(self.num_envs, dtype=bool))
+            
+            obs, info = self.env.reset(**kwargs)
+            self._setup_callbacks()
+            
+            for i in np.where(reset_mask)[0]:
+                if self.steps[i] == 0 and self.record:
+                    self.video_frames[i] = []
+                self.max_reward[i] = 0
+                self.steps[i] = 0
+                self.episode_info[i] = {}
+            
+            self._callbacks_paused = False
+        else:
+            if self.steps > 0:
+                if self.csv_writer:
+                    self._log_episode()
+                if self.record:
+                    self._save_video()
+                self.episode_idx += 1
+            obs, info = self.env.reset(**kwargs)
+            self._setup_callbacks()
+            
+            self.max_reward = 0
+            self.steps = 0
+            self.episode_info = {}
+            if self.record:
+                self.video_frames = []
+            self._callbacks_paused = False
+        return obs, info
+    
+    def step(self, action):
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        
+        if self.render_freq == 0:
+            if self.is_vector:
+                # AsyncVectorEnv doesn't provide direct access to individual envs
+                # Recording from observations if available
+                if self.record and 'rgb_ego' in obs and 'rgb_exo' in obs:
+                    for i in range(self.num_envs):
+                        frame_ego = obs['rgb_ego'][i]
+                        frame_exo = obs['rgb_exo'][i]
+                        frame = np.concatenate((frame_ego, frame_exo), axis=1)
+                        frame = self._process_frame(frame)
+                        self.video_frames[i].append(frame)
+            else:
+                actual_env = self.env.unwrapped
+                if actual_env._should_render:
+                    actual_env.render()
+                if self.record and 'rgb_ego' in obs and 'rgb_exo' in obs:
+                    frame_ego = obs['rgb_ego']
+                    frame_exo = obs['rgb_exo']
+                    frame = np.concatenate((frame_ego, frame_exo), axis=1)
+                    frame = self._process_frame(frame)
+                    self.video_frames.append(frame)
+        
+        if self.is_vector:
+            self.max_reward = np.maximum(self.max_reward, reward)
+            self.steps += 1
+            for i in range(self.num_envs):
+                if isinstance(info, dict):
+                    env_info = {}
+                    for k, v in info.items():
+                        if k.startswith('_'):
+                            continue
+                        if isinstance(v, (list, np.ndarray)) and len(v) > i:
+                            env_info[k] = v[i]
+                        else:
+                            env_info[k] = v
+                else:
+                    env_info = info[i] if isinstance(info, (list, tuple)) else {}
+                self.episode_info[i].update(env_info)
+            
+            dones = np.logical_or(terminated, truncated)
+            if isinstance(dones, bool):
+                dones = np.array([dones] * self.num_envs)
+        else:
+            self.max_reward = max(self.max_reward, reward)
+            self.steps += 1
+            self.episode_info.update(info)
+            info['episode_stats'] = {'episode': self.episode_idx, 'max_reward': self.max_reward, 'steps': self.steps}
+        
+        return obs, reward, terminated, truncated, info
+    
+    def _log_episode(self, env_idx=None):
+        if env_idx is not None:
+            log_data = {"episode": self.global_episode_counter, "max_reward": float(self.max_reward[env_idx]), "steps": int(self.steps[env_idx]), **self.episode_info[env_idx]}
+            self.global_episode_counter += 1
+        else:
+            log_data = {"episode": self.episode_idx, "max_reward": self.max_reward, "steps": self.steps, **self.episode_info}
+        
+        filtered_data = {}
+        for col in self.columns:
+            value = log_data.get(col, "")
+            if isinstance(value, np.ndarray):
+                value = value.flatten().tolist()
+            filtered_data[col] = value
+        
+        self.csv_writer.writerow(filtered_data)
+        self.csv_file.flush()
+    
+    def log_episodes(self, env_indices):
+        if self.is_vector:
+            for i in env_indices:
+                if self.csv_writer:
+                    self._log_episode(i)
+                if self.record:
+                    self._save_video(i)
+                    self.video_frames[i] = []
+                self.episode_idx[i] += 1
+    
+    def _save_video(self, env_idx=None):
+        if env_idx is not None:
+            frames = self.video_frames[env_idx]
+            episode_num = int(self.episode_idx[env_idx])
+            video_filename = f"env_{env_idx}_episode_{episode_num}.mp4"
+        else:
+            frames = self.video_frames
+            episode_num = self.episode_idx
+            video_filename = f"episode_{episode_num}.mp4"
+        
+        if not frames:
+            return
+        
+        video_path = os.path.join(self.video_dir, video_filename)
+        height, width = frames[0].shape[:2]
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(video_path, fourcc, 30, (width, height))
+        
+        for frame in frames:
+            out.write(frame)
+        
+        out.release()
+    
+    def close(self):
+        if self.is_vector:
+            for i in range(self.num_envs):
+                if self.steps[i] > 0:
+                    if self.csv_writer:
+                        self._log_episode(i)
+                    if self.record:
+                        self._save_video(i)
+        else:
+            if self.steps > 0:
+                if self.csv_writer:
+                    self._log_episode()
+                if self.record:
+                    self._save_video()
+        if self.csv_file:
+            self.csv_file.close()
+        self.env.close()
+    
+    def __getattr__(self, name):
+        return getattr(self.env, name)

egogym/wrappers/unprivileged_chatgpt.py

@@ -0,0 +1,163 @@
+import os
+import numpy as np
+import time
+import re
+import io
+import json
+import base64
+from PIL import Image
+from egogym.utils import pixel_to_world
+
+MODEL = "gpt-4o"
+PROMPT = """Get all points matching the following object: {object_name}. The label returned should be an identifying name for the object detected.
+The answer should follow the json format: [{{"point": [y, x], "label": "{object_name}"}}, ...]. The points are in [y, x] format normalized to 0-1000."""
+
+class UnprivilegedChatGPT:
+    
+    def __init__(self, env, api_key=None):
+        try:
+            from openai import OpenAI
+        except ImportError:
+            raise ImportError("openai package is not installed. Please install it with `pip install openai`.")
+
+        self.env = env
+        self.unprivileged_T_world_object = None
+        
+        if api_key is None:
+            api_key = os.environ.get("OPENAI_API_KEY")
+    
+        
+        self.client = OpenAI(api_key=api_key)
+    
+    def _call_api_with_retry(self, messages, max_retries=3, base_delay=2):
+        """Call OpenAI API with retry logic for transient errors."""
+        for attempt in range(max_retries):
+            try:
+                return self.client.chat.completions.create(
+                    model=MODEL,
+                    messages=messages,
+                    temperature=0.5,
+                    max_tokens=1024
+                )
+            except Exception as e:
+                error_code = getattr(e, 'status_code', None)
+                if error_code in [503, 429] and attempt < max_retries - 1:
+                    delay = base_delay * (2 ** attempt)
+                    print(f"API error {error_code}: {str(e)}. Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})")
+                    time.sleep(delay)
+                    continue
+                else:
+                    raise
+    
+    def _parse_json_point_yx(self, text):
+        """Parse JSON format with point coordinates in [y, x] format normalized to 0-1000."""
+        try:
+            json_match = re.search(r'```(?:json)?\s*(\[.*?\])\s*```', text, re.DOTALL)
+            if json_match:
+                json_str = json_match.group(1)
+            else:
+                json_match = re.search(r'(\[.*?\])', text, re.DOTALL)
+                if json_match:
+                    json_str = json_match.group(1)
+                else:
+                    return None
+            
+            data = json.loads(json_str)
+            if not isinstance(data, list) or len(data) == 0:
+                return None
+            
+            point_data = data[0]
+            if 'point' in point_data:
+                point = point_data['point']
+                if isinstance(point, list) and len(point) >= 2:
+                    return (float(point[0]), float(point[1]))
+            
+            return None
+            
+        except (json.JSONDecodeError, KeyError, ValueError, IndexError) as e:
+            return None
+    
+    def _infer_point_sync(self, rgb: np.ndarray, object_name: str) -> np.ndarray:
+        """Use OpenAI GPT-4o 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)
+        img_width, img_height = image.size
+        
+        prompt = PROMPT.format(object_name=object_name)
+        
+        try:
+            buf = io.BytesIO()
+            image.save(buf, format='PNG')
+            image_bytes = buf.getvalue()
+            base64_image = base64.b64encode(image_bytes).decode('utf-8')
+            
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/png;base64,{base64_image}"
+                            }
+                        },
+                        {
+                            "type": "text",
+                            "text": prompt
+                        }
+                    ]
+                }
+            ]
+            
+            response = self._call_api_with_retry(messages)
+            generated_text = response.choices[0].message.content
+            
+            json_point = self._parse_json_point_yx(generated_text)
+            if json_point:
+                y_pos, x_pos = json_point
+                x_norm = float(x_pos) / 1000.0
+                y_norm = float(y_pos) / 1000.0
+                
+                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:
+            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/wrappers/unprivileged_gemini.py

@@ -0,0 +1,157 @@
+import os
+import numpy as np
+import time
+import re
+import io
+import json
+from PIL import Image
+from egogym.utils import pixel_to_world
+
+MODEL = "gemini-robotics-er-1.5-preview"
+PROMPT = """Get all points matching the following object: {object_name}. The label returned should be an identifying name for the object detected.
+The answer should follow the json format: [{{"point": [y, x], "label": "{object_name}"}}, ...]. The points are in [y, x] format normalized to 0-1000."""
+
+class UnprivilegedGemini:
+    
+    def __init__(self, env, api_key=None):
+        try:
+            import google.genai
+            from google.genai import errors, types
+        except ImportError:
+            raise ImportError("google-genai package is not installed. Please install it with `pip install google-genai`.")
+        
+
+        self.env = env
+        self.unprivileged_T_world_object = None
+        
+        if api_key is None:
+            api_key = os.getenv("GEMINI_API_KEY", "AIzaSyBQmaK2VE3y8qxOGOEliQsveTDpfK2LyLc")
+        
+        self.client = genai.Client(api_key=api_key)
+    
+    def _call_api_with_retry(self, model, contents, config, max_retries=3, base_delay=2):
+        """Call Gemini API with retry logic for transient errors."""
+        for attempt in range(max_retries):
+            try:
+                return self.client.models.generate_content(
+                    model=model,
+                    contents=contents,
+                    config=config
+                )
+            except errors.ServerError as e:
+                if e.code in [503, 429] and attempt < max_retries - 1:
+                    delay = base_delay * (2 ** attempt)
+                    error_msg = e.message if e.message else "Service unavailable"
+                    print(f"API error {e.code}: {error_msg}. Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})")
+                    time.sleep(delay)
+                    continue
+                else:
+                    raise
+            except Exception as e:
+                raise
+    
+    def _parse_json_point_yx(self, text):
+        """Parse JSON format with point coordinates in [y, x] format normalized to 0-1000."""
+        try:
+            json_match = re.search(r'```(?:json)?\s*(\[.*?\])\s*```', text, re.DOTALL)
+            if json_match:
+                json_str = json_match.group(1)
+            else:
+                json_match = re.search(r'(\[.*?\])', text, re.DOTALL)
+                if json_match:
+                    json_str = json_match.group(1)
+                else:
+                    return None
+            
+            data = json.loads(json_str)
+            if not isinstance(data, list) or len(data) == 0:
+                return None
+            
+            point_data = data[0]
+            if 'point' in point_data:
+                point = point_data['point']
+                if isinstance(point, list) and len(point) >= 2:
+                    return (float(point[0]), float(point[1]))
+            
+            return None
+            
+        except (json.JSONDecodeError, KeyError, ValueError, IndexError) as e:
+            return None
+    
+    def _infer_point_sync(self, rgb: np.ndarray, object_name: str) -> np.ndarray:
+        """Use Gemini 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)
+        img_width, img_height = image.size
+        
+        prompt = PROMPT.format(object_name=object_name)
+        
+        try:
+            buf = io.BytesIO()
+            image.save(buf, format='PNG')
+            image_bytes = buf.getvalue()
+            
+            contents = [
+                types.Part.from_bytes(
+                    data=image_bytes,
+                    mime_type='image/png',
+                ),
+                prompt
+            ]
+            
+            config = types.GenerateContentConfig(
+                temperature=0.5,
+                thinking_config=types.ThinkingConfig(thinking_budget=0)
+            )
+            response = self._call_api_with_retry(MODEL, contents, config)
+            generated_text = response.text
+            
+            json_point = self._parse_json_point_yx(generated_text)
+            if json_point:
+                y_pos, x_pos = json_point
+                x_norm = float(x_pos) / 1000.0
+                y_norm = float(y_pos) / 1000.0
+                
+                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:
+            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)