egogym 0.1.0__py3-none-any.whl

egogym/evaluate.py

@@ -0,0 +1,216 @@
+import os
+import time
+import cv2
+import gymnasium as gym
+import numpy as np
+from tqdm import tqdm
+from itertools import cycle, islice
+import multiprocessing
+import json
+
+import egogym
+from egogym.wrappers import EpisodeMonitor, get_vlm_wrapper
+import egogym.assets.constants as constants
+
+def evaluate_policy(
+        task_name="pick",
+        policy=None,
+        robot="rum",
+        action_space="delta",
+        num_objs=1,
+        num_episodes=100,
+        max_steps=50,
+        render=False,
+        record=False,
+        render_freq=0,
+        render_size=(299, 224),
+        num_envs=1,
+        seed=42,
+        log_file=None,
+        objects_set=None,
+        reward_threshold=None,
+        use_unprivileged_vlm=None,
+        logs_dir="logs"
+    ):
+
+    num_envs = max(1, min(num_envs, num_episodes))
+    if num_envs > 1 and render_freq > 0:
+        print(f"Warning: render_freq={render_freq} is not supported with multiple environments. Setting render_freq=0.")
+        render_freq = 0
+    task_name = task_name.lower()
+    if objects_set is None:
+        if task_name == "pick":
+            objects_set = constants.lite_pick_objects_set
+        elif task_name == "open":
+            objects_set = constants.all_open_objects_set
+    else:
+        objects_set_options = {
+            'all_pick': constants.all_pick_objects_set,
+            'lite_pick': constants.lite_pick_objects_set,
+            'diverse_pick': constants.diverse_pick_objects_set,
+            'full_eval': constants.full_eval_objects_set,
+            'all_open': constants.all_open_objects_set,
+            "all_close": constants.all_close_objects_set,
+            'cabinet': constants.cabinet_objects_set,
+            'drawer': constants.drawer_objects_set,
+        }
+        if objects_set in objects_set_options:
+            objects_set = objects_set_options[objects_set]
+        else:
+            raise ValueError(f"Unknown objects_set: {objects_set}")
+        
+    if reward_threshold is None:
+        if task_name == "pick":
+            reward_threshold = 0.05
+        elif task_name == "open":
+            reward_threshold = 0.5
+
+    if not log_file:
+        log_file = f"evaluation_{int(time.time())}"
+
+    logs_dir = f"{logs_dir}/{log_file}"
+    os.makedirs(logs_dir, exist_ok=True)
+    
+    eval_args = {
+        "task_name": task_name,
+        "robot": robot,
+        "action_space": action_space,
+        "num_objs": num_objs,
+        "num_episodes": num_episodes,
+        "max_steps": max_steps,
+        "render": render,
+        "record": record,
+        "render_freq": render_freq,
+        "render_size": render_size,
+        "num_envs": num_envs,
+        "seed": seed,
+        "log_file": log_file,
+        "objects_set": objects_set,
+        "reward_threshold": reward_threshold,
+    }
+    
+    with open(os.path.join(logs_dir, "eval_args.json"), "w") as f:
+        json.dump(eval_args, f, indent=4)
+
+    if num_envs > 1:
+        multiprocessing.set_start_method("spawn", force=True)
+        
+        expanded_objects = list(islice(cycle(objects_set), num_episodes))
+        env_objects = [expanded_objects[i::num_envs] for i in range(num_envs)]
+        
+        def make_env(env_id):
+            def _init():
+                env = gym.make(
+                    f"Egogym-{task_name.capitalize()}-v0",
+                    robot=robot,
+                    action_space=action_space,
+                    render_mode='human' if render else None,
+                    render_size=render_size,
+                    num_objs=num_objs,
+                    seed=seed + env_id,
+                    objects_set=env_objects[env_id]
+                )
+                if use_unprivileged_vlm is not None:
+                    wrapper_cls = get_vlm_wrapper(use_unprivileged_vlm)
+                    env = wrapper_cls(env)
+                return env
+            return _init
+        
+        env = gym.vector.AsyncVectorEnv(
+            [make_env(i) for i in range(num_envs)],
+            autoreset_mode=gym.vector.AutoresetMode.DISABLED
+        )
+        env = EpisodeMonitor(env, logs_dir=logs_dir, record=record, render_freq=render_freq, num_envs=num_envs)
+        obs, _ = env.reset()
+        env_episode_count = np.zeros(num_envs, dtype=int)
+        env_step_count = np.zeros(num_envs, dtype=int)
+        env_success_count = 0
+        episodes_per_env = np.array([len(env_objects[i]) for i in range(num_envs)])
+        
+        pbar = tqdm(total=num_episodes, desc="Evaluating | Success Rate: 0.0%")
+        
+        while np.any(env_episode_count < episodes_per_env):
+            actions = policy.get_action(obs)
+            inactive = env_episode_count >= episodes_per_env
+            actions[inactive] = None
+
+            obs, rewards, terminated, truncated, info = env.step(actions)
+            env_step_count += 1
+            
+            if render:
+                ego_exo = np.concatenate([obs["rgb_ego"], obs["rgb_exo"]], axis=2)
+                combined_view = np.concatenate(ego_exo, axis=0)
+                cv2.imshow("Ego View", cv2.cvtColor(combined_view, cv2.COLOR_RGB2BGR))
+                cv2.waitKey(1)
+            
+            dones = (rewards > reward_threshold) | terminated | (env_step_count >= max_steps)
+            done_indices = np.where(dones)[0]
+
+            active_done = done_indices[env_episode_count[done_indices] < episodes_per_env[done_indices]]
+            
+            if len(active_done) > 0:
+                env.log_episodes(active_done)
+                
+                successes = rewards[active_done] > reward_threshold
+                env_success_count += np.sum(successes)
+                
+                env_episode_count[active_done] += 1
+                pbar.update(len(active_done))
+                
+                total_completed = np.sum(env_episode_count)
+                success_rate = (env_success_count / total_completed) * 100 if total_completed > 0 else 0
+                pbar.set_description(f"Evaluating | Success Rate: {success_rate:.1f}%")
+                
+                policy.reset(active_done.tolist())
+                
+                reset_mask = np.zeros(num_envs, dtype=bool)
+                reset_mask[active_done] = True
+                obs, _ = env.reset(options={"reset_mask": reset_mask})
+                env_step_count[active_done] = 0
+            
+        pbar.close()
+        success_rate = (env_success_count / num_episodes) * 100
+        return success_rate
+        
+    else:
+
+        env = gym.make(
+            f"Egogym-{task_name.capitalize()}-v0",
+            robot=robot,
+            action_space=action_space,
+            render_mode='human' if render else None,
+            render_size=render_size,
+            num_objs=num_objs,
+            seed=seed,
+            objects_set=objects_set
+        )
+        if use_unprivileged_vlm is not None:
+            wrapper_cls = get_vlm_wrapper(use_unprivileged_vlm)
+            env = wrapper_cls(env)
+        env = EpisodeMonitor(env, logs_dir=logs_dir, record=record, render_freq=render_freq, num_envs=num_envs)
+        
+        success_count = 0
+        pbar = tqdm(range(num_episodes), desc="Evaluating | Success Rate: 0.0%")
+        for episode_idx in pbar:
+            obs, _ = env.reset()
+            policy.reset()
+            
+            success = False
+            for _ in range(max_steps):
+                action = policy.get_action(obs)
+                obs, reward, terminated, truncated, _ = env.step(action[0])
+                    
+                if reward > reward_threshold:
+                    success = True
+                    break
+                if terminated or truncated:
+                    break
+            
+            if success:
+                success_count += 1
+            
+            success_rate = (success_count / (episode_idx + 1)) * 100
+            pbar.set_description(f"Evaluating | Success Rate: {success_rate:.1f}%")
+
+    env.close()
+    return (success_count / num_episodes) * 100 if num_envs == 1 else success_rate

egogym/managers/__init__.py

@@ -0,0 +1,2 @@
+from .objects_managers import ObjectsManager
+from .textures_manager import TexturesManager

egogym/managers/objects_managers.py

@@ -0,0 +1,30 @@
+from egogym.utils import make_cabinet_xml, make_drawer_xml
+from egogym.components import Object, ArticulableObject
+
+class ObjectsManager:
+    def __init__(self, objects_dict, np_random, shuffle=True):
+        self.np_random = np_random
+        sorted_objects = sorted(objects_dict, key=lambda x: x[0])
+        if shuffle:
+            np_random.shuffle(sorted_objects)
+        self.objects_dict = sorted_objects
+        self.index = -1
+
+    def sample(self, random=False):
+        if random:
+            index = self.np_random.integers(len(self.objects_dict))
+        else:
+            self.index = (self.index + 1) % len(self.objects_dict)
+            index = self.index
+        if self.objects_dict[index][1] == "procedurally_generated":
+            return ArticulableObject(
+                name=self.objects_dict[index][0],
+                func=globals()[f"make_{self.objects_dict[index][0]}_xml"],
+                np_random=self.np_random,
+            )
+        else:
+            return Object(
+                name=self.objects_dict[index][0],
+                path=self.objects_dict[index][1],
+                np_random=self.np_random,
+            )

egogym/managers/textures_manager.py

@@ -0,0 +1,21 @@
+import numpy as np
+
+from egogym.components import Object
+
+
+class TexturesManager:
+    def __init__(self, textures_list, np_random):
+        self.textures_list = textures_list
+        self.np_random = np_random
+        sorted_textures = sorted(textures_list)
+        self.np_random.shuffle(sorted_textures)
+        self.textures_list = sorted_textures
+        self.index = -1
+
+    def sample(self, random=False):
+        if random:
+            index = self.np_random.integers(len(self.textures_list))
+        else:
+            self.index = (self.index + 1) % len(self.textures_list)
+            index = self.index
+        return self.textures_list[index]

egogym/misc/molmo_client.py

@@ -0,0 +1,49 @@
+import json
+import numpy as np
+import websockets
+
+class MolmoClient:
+    def __init__(self, host="localhost", port=8765):
+        self.uri = f"ws://{host}:{port}"
+        self.websocket = None
+
+    async def connect(self):
+        self.websocket = await websockets.connect(
+            self.uri,
+            max_size=50 * 1024 * 1024,
+        )
+
+    async def disconnect(self):
+        if self.websocket:
+            await self.websocket.close()
+
+    async def infer_point(self, rgb: np.ndarray, object_name: str = None, prompt: str = None):
+        if not self.websocket:
+            raise RuntimeError("Not connected to server. Call connect() first.")
+
+        req = {
+            "action": "infer_point",
+            "rgb": rgb.tolist(),
+        }
+
+        if object_name:
+            req["object_name"] = object_name
+        elif prompt:
+            req["prompt"] = prompt
+        else:
+            raise ValueError("Provide either 'object_name' or 'prompt'")
+
+        await self.websocket.send(json.dumps(req))
+        resp = json.loads(await self.websocket.recv())
+
+        if resp["status"] == "error":
+            raise RuntimeError(f"Server error: {resp['message']}")
+
+        return np.array(resp["point"], dtype=np.float32)
+
+    async def __aenter__(self):
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.disconnect()

egogym/misc/molmo_server.py

@@ -0,0 +1,197 @@
+import argparse
+import asyncio
+import json
+import logging
+import re
+import numpy as np
+import torch
+import websockets
+from PIL import Image
+
+logging.basicConfig(level=logging.INFO)
+log = logging.getLogger(__name__)
+
+
+def extract_molmo_points(molmo_output: str):
+    points = []
+    for match in re.finditer(
+        r'x\d*="\s*([0-9]+(?:\.[0-9]+)?)"\s+y\d*="\s*([0-9]+(?:\.[0-9]+)?)"',
+        molmo_output,
+    ):
+        try:
+            p = np.array([float(match.group(1)), float(match.group(2))], dtype=np.float32)
+        except ValueError:
+            continue
+
+        if np.max(p) > 100:
+            continue
+
+        points.append(p / 100.0)
+
+    return points
+
+
+class Molmo:
+    def __init__(self, model_name="allenai/Molmo-7B-D-0924"):
+        try:
+            from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
+        except ImportError:
+            raise ImportError("transformers package is not installed. Please install it with `pip install transformers`.")
+        
+        self.processor = AutoProcessor.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            torch_dtype="auto",
+            device_map="auto",
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            torch_dtype="auto",
+            device_map="auto",
+        )
+
+    def infer(self, rgb: np.ndarray, prompt: str) -> str:
+        image = Image.fromarray(rgb)
+        inputs = self.processor.process(images=[image], text=prompt)
+        inputs = {k: v.to(self.model.device).unsqueeze(0) for k, v in inputs.items()}
+
+        with torch.inference_mode():
+            output = self.model.generate_from_batch(
+                inputs,
+                GenerationConfig(max_new_tokens=200, stop_strings="<|endoftext|>"),
+                tokenizer=self.processor.tokenizer,
+            )
+
+        gen_tokens = output[0, inputs["input_ids"].size(1):]
+        return self.processor.tokenizer.decode(gen_tokens, skip_special_tokens=True)
+
+    def infer_point(self, rgb: np.ndarray, prompt: str) -> np.ndarray:
+        text = self.infer(rgb, prompt)
+        log.info(f"Molmo output: {text}")
+        points = extract_molmo_points(text)
+        return points[0] if points else np.array([0.5, 0.5], dtype=np.float32)
+
+
+class MolmoClient:
+    def __init__(self, host="localhost", port=8765):
+        self.uri = f"ws://{host}:{port}"
+        self.websocket = None
+
+    async def connect(self):
+        self.websocket = await websockets.connect(
+            self.uri,
+            max_size=50 * 1024 * 1024,
+        )
+        log.info(f"Connected to Molmo server at {self.uri}")
+
+    async def disconnect(self):
+        if self.websocket:
+            await self.websocket.close()
+            log.info("Disconnected from Molmo server")
+
+    async def infer_point(self, rgb: np.ndarray, object_name: str = None, prompt: str = None):
+        if not self.websocket:
+            raise RuntimeError("Not connected to server. Call connect() first.")
+
+        req = {
+            "action": "infer_point",
+            "rgb": rgb.tolist(),
+        }
+
+        if object_name:
+            req["object_name"] = object_name
+        elif prompt:
+            req["prompt"] = prompt
+        else:
+            raise ValueError("Provide either 'object_name' or 'prompt'")
+
+        await self.websocket.send(json.dumps(req))
+        resp = json.loads(await self.websocket.recv())
+
+        if resp["status"] == "error":
+            raise RuntimeError(f"Server error: {resp['message']}")
+
+        return np.array(resp["point"], dtype=np.float32)
+
+    async def __aenter__(self):
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.disconnect()
+
+
+class MolmoWebSocketServer:
+    def __init__(self, host="0.0.0.0", port=8765):
+        self.host = host
+        self.port = port
+        self.molmo = Molmo()
+
+    async def handle_client(self, websocket):
+        client_id = id(websocket)
+        log.info(f"Client connected: {client_id}")
+
+        try:
+            async for msg in websocket:
+                try:
+                    req = json.loads(msg)
+                    action = req.get("action")
+
+                    if action != "infer_point":
+                        raise ValueError("Only 'infer_point' action is supported")
+
+                    rgb = np.array(req["rgb"], dtype=np.uint8)
+
+                    if "object_name" in req:
+                        prompt = f"Point to the center of the {req['object_name']}."
+                        label = req["object_name"].replace(" ", "_")
+                    elif "prompt" in req:
+                        prompt = req["prompt"]
+                        label = "custom_prompt"
+                    else:
+                        raise ValueError("Provide either 'object_name' or 'prompt'")
+
+                    point = self.molmo.infer_point(rgb, prompt)
+
+                    resp = {
+                        "status": "ok",
+                        "action": "infer_point",
+                        "point": point.tolist(),
+                    }
+
+                except Exception as e:
+                    log.exception("Request error")
+                    resp = {"status": "error", "message": str(e)}
+
+                await websocket.send(json.dumps(resp))
+
+        except websockets.exceptions.ConnectionClosed:
+            log.info(f"Client disconnected: {client_id}")
+
+    async def start(self):
+        log.info(f"Starting Molmo server on ws://{self.host}:{self.port}")
+        async with websockets.serve(
+            self.handle_client,
+            self.host,
+            self.port,
+            max_size=50 * 1024 * 1024,
+        ):
+            await asyncio.Future()
+
+
+def main():
+    parser = argparse.ArgumentParser("Molmo Pointing Server")
+    parser.add_argument("--host", default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=8765)
+    args = parser.parse_args()
+
+    server = MolmoWebSocketServer(
+        host=args.host,
+        port=args.port,
+    )
+    asyncio.run(server.start())
+
+
+if __name__ == "__main__":
+    main()

egogym/policies/__init__.py

@@ -0,0 +1 @@
+from .base_policy import BasePolicy

egogym/policies/base_policy.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+
+class BasePolicy(ABC):
+    
+    def __init__(self):
+        self.name = "base"
+        
+    @abstractmethod
+    def get_action(self, obs):
+        pass
+    
+    def reset(self):
+        pass