Upload server/robosim/sim_wrapper.py with huggingface_hub

server/robosim/sim_wrapper.py

@@ -0,0 +1,358 @@
+"""
+SimWrapper — physics backend for RoboReplan.
+
+Two modes:
+  use_stub=False  Real MuJoCo + robosuite PickPlace environment.
+                  Object positions, blocking, and grasp success come from
+                  actual physics. This is what makes training meaningful.
+
+  use_stub=True   Lightweight Python sim for fast local testing.
+                  Same interface, no physics dependency.
+
+The wrapper always exposes the same symbolic SimState so the planning
+layer above never needs to know which backend is running.
+"""
+import random
+import numpy as np
+from dataclasses import dataclass, field
+from typing import Optional
+
+from .randomizer import randomize_scenario, ScenarioConfig
+
+
+# ── Symbolic state types ───────────────────────────────────────────────
+
+@dataclass
+class ObjectState:
+    name: str
+    pos: np.ndarray
+    reachable: bool = True
+    blocking: Optional[str] = None
+    in_bin: Optional[str] = None
+    is_held: bool = False
+
+
+@dataclass
+class SimState:
+    objects: dict = field(default_factory=dict)
+    gripper_pos: np.ndarray = field(default_factory=lambda: np.zeros(3))
+    gripper_open: bool = True
+    holding: Optional[str] = None
+
+
+# ── Main wrapper ───────────────────────────────────────────────────────
+
+class SimWrapper:
+    """
+    Wraps either robosuite (real) or a Python stub (fast testing).
+    Always produces symbolic SimState for the planning layer.
+    """
+
+    def __init__(self, use_stub: bool = True):
+        self.use_stub = use_stub
+        self._env = None
+        self._controller = None
+        self._state = SimState()
+        self._last_moved_to: Optional[str] = None
+        self._current_cfg: Optional[ScenarioConfig] = None
+        self._facing: str = "N"
+
+        if not use_stub:
+            self._init_robosuite()
+
+    # ── Init ───────────────────────────────────────────────────────────
+
+    def _init_robosuite(self):
+        """
+        Initialize a real robosuite PickPlace environment.
+
+        We use PickPlace because it already has:
+        - Panda arm (most standard research robot)
+        - Multiple objects with real physics
+        - Multiple bin targets
+        - OSC_POSE controller (operational space — moves in Cartesian coords,
+          which our high-level controller can use directly)
+        """
+        try:
+            import robosuite as suite
+            from robosuite.controllers import load_composite_controller_config
+
+            controller_config = load_composite_controller_config(controller="BASIC")
+
+            self._env = suite.make(
+                env_name="PickPlace",
+                robots="Panda",
+                controller_configs=controller_config,
+                has_renderer=False,           # no display on server
+                has_offscreen_renderer=True,  # needed for camera obs
+                use_camera_obs=True,
+                camera_names=["frontview", "agentview"],
+                camera_heights=128,
+                camera_widths=128,
+                reward_shaping=False,         # we compute our own reward
+                control_freq=20,
+                single_object_mode=0,         # all objects
+                object_type=None,             # random objects
+            )
+
+            from .controller import MotionController
+            self._controller = MotionController(self._env)
+            print("[SimWrapper] robosuite PickPlace loaded (Panda arm)")
+
+        except ImportError:
+            print("[SimWrapper] robosuite not installed → falling back to stub")
+            self.use_stub = True
+        except Exception as e:
+            print(f"[SimWrapper] robosuite init failed: {e} → falling back to stub")
+            self.use_stub = True
+
+    # ── Reset ──────────────────────────────────────────────────────────
+
+    def reset(self, scenario: str = "random") -> tuple[SimState, ScenarioConfig]:
+        """Reset scene. Returns (SimState, ScenarioConfig)."""
+        force_blocked = random.random() < 0.6
+        cfg = randomize_scenario(force_blocked=force_blocked)
+
+        if not self.use_stub and self._env is not None:
+            self._reset_robosuite(cfg)
+        else:
+            self._build_state_from_config(cfg)
+
+        return self._state, cfg
+
+    def _reset_robosuite(self, cfg: ScenarioConfig):
+        """Reset robosuite and sync symbolic state from physics."""
+        obs = self._env.reset()
+        self._sync_state_from_obs(obs, cfg)
+
+    def _sync_state_from_obs(self, obs: dict, cfg: ScenarioConfig):
+        """
+        Extract symbolic state from robosuite observation dict.
+        Uses the perception layer to detect blocking from real 3D positions.
+        """
+        try:
+            from .perception import extract_scene
+            scene = extract_scene(
+                mj_data=self._env.sim.data,
+                mj_model=self._env.sim.model,
+                robot_name="robot0",
+                object_names=list(cfg.objects),
+            )
+            objects = {}
+            for name, p in scene.objects.items():
+                objects[name] = ObjectState(
+                    name=name,
+                    pos=p.pos,
+                    reachable=p.reachable,
+                    blocking=p.blocking,
+                    in_bin=p.in_bin,
+                    is_held=p.is_held,
+                )
+            self._state = SimState(
+                objects=objects,
+                gripper_pos=scene.gripper_pos,
+                gripper_open=scene.gripper_open,
+                holding=scene.holding,
+            )
+        except Exception as e:
+            print(f"[SimWrapper] perception sync failed: {e}, using stub fallback")
+            self._build_state_from_config(cfg)
+
+        self._current_cfg = cfg
+
+    # ── Stub state builder ─────────────────────────────────────────────
+
+    def get_last_moved_to(self) -> Optional[str]:
+        return self._last_moved_to
+
+    def _build_state_from_config(self, cfg: ScenarioConfig):
+        """Build stub SimState from randomized scenario config."""
+        self._last_moved_to = None
+        self._facing = "N"
+        objects = {}
+        for obj_name in cfg.objects:
+            x, y = cfg.positions.get(obj_name, (0.0, 0.0))
+            is_blocked = obj_name in cfg.blockers.values()
+            objects[obj_name] = ObjectState(
+                name=obj_name,
+                pos=np.array([x, y, 0.82]),
+                reachable=not is_blocked,
+                blocking=cfg.blockers.get(obj_name),
+            )
+        self._state = SimState(
+            objects=objects,
+            gripper_pos=np.array([0.0, 0.25, 1.0]),
+            gripper_open=True,
+            holding=None,
+        )
+        self._current_cfg = cfg
+
+    def get_facing(self) -> str:
+        return self._facing
+
+    def _cell_from_pos(self, pos: np.ndarray) -> tuple[int, int]:
+        x = int(round(float(pos[0]) / 0.1))
+        y = int(round(float(pos[1]) / 0.1))
+        return max(-3, min(3, x)), max(-3, min(3, y))
+
+    def _step_gripper(self, dx: int, dy: int) -> None:
+        s = self._state
+        cell_x, cell_y = self._cell_from_pos(s.gripper_pos)
+        nx = max(-3, min(3, cell_x + dx))
+        ny = max(-3, min(3, cell_y + dy))
+        s.gripper_pos = np.array([nx * 0.1, ny * 0.1, s.gripper_pos[2]])
+
+    def _rotate(self, clockwise: bool) -> None:
+        dirs = ["N", "E", "S", "W"]
+        idx = dirs.index(self._facing)
+        self._facing = dirs[(idx + (1 if clockwise else -1)) % 4]
+
+    def _is_adjacent(self, obj: ObjectState) -> bool:
+        gx, gy = self._cell_from_pos(self._state.gripper_pos)
+        ox, oy = self._cell_from_pos(obj.pos)
+        return abs(gx - ox) + abs(gy - oy) <= 1
+
+    # ── Execute action ─────────────────────────────────────────────────
+
+    def execute(self, action: str) -> str:
+        """Execute a high-level action. Returns result string."""
+        if not self.use_stub and self._env is not None and self._controller is not None:
+            return self._execute_real(action)
+        return self._execute_stub(action)
+
+    def _execute_real(self, action: str) -> str:
+        """Execute via real robosuite physics + motion controller."""
+        result = self._controller.execute(action)
+        # Re-sync symbolic state from physics
+        obs = self._env._get_observations()
+        if self._current_cfg:
+            self._sync_state_from_obs(obs, self._current_cfg)
+        return result
+
+    def _execute_stub(self, action: str) -> str:
+        """Execute in the lightweight Python stub."""
+        s = self._state
+
+        if action == "SCAN_SCENE":
+            return "SUCCESS"
+        elif action == "MOVE_NORTH":
+            self._step_gripper(0, 1)
+            return "SUCCESS"
+        elif action == "MOVE_SOUTH":
+            self._step_gripper(0, -1)
+            return "SUCCESS"
+        elif action == "MOVE_EAST":
+            self._step_gripper(1, 0)
+            return "SUCCESS"
+        elif action == "MOVE_WEST":
+            self._step_gripper(-1, 0)
+            return "SUCCESS"
+        elif action == "ROTATE_LEFT":
+            self._rotate(clockwise=False)
+            return "SUCCESS"
+        elif action == "ROTATE_RIGHT":
+            self._rotate(clockwise=True)
+            return "SUCCESS"
+
+        elif action.startswith("MOVE_TO_"):
+            color = action[len("MOVE_TO_"):].lower()
+            # Try "<color>_block" first (default pack), then bare name (professional packs)
+            name = color + "_block" if (color + "_block") in s.objects else color
+            if name not in s.objects:
+                return "FAILED_INVALID"
+            obj = s.objects[name]
+            if not obj.reachable:
+                return "FAILED_BLOCKED"
+            s.gripper_pos = obj.pos.copy() + np.array([0, 0, 0.05])
+            self._last_moved_to = name
+            return "SUCCESS"
+
+        elif action == "PICK":
+            if s.holding is not None:
+                return "FAILED_INVALID"
+            candidates = []
+            for obj in s.objects.values():
+                if obj.reachable and not obj.is_held and obj.in_bin is None:
+                    dist = np.linalg.norm(s.gripper_pos[:2] - obj.pos[:2])
+                    candidates.append((dist, obj))
+            # Prefer the object we last moved to
+            candidates.sort(key=lambda x: (
+                0 if x[1].name == self._last_moved_to else 1, x[0]
+            ))
+            for _, obj in candidates:
+                dist = np.linalg.norm(s.gripper_pos[:2] - obj.pos[:2])
+                if dist < 0.25:
+                    obj.is_held = True
+                    s.holding = obj.name
+                    s.gripper_open = False
+                    self._last_moved_to = None
+                    return "SUCCESS"
+            return "FAILED_EMPTY"
+
+        elif action in ("PLACE_BIN_A", "PLACE_BIN_B"):
+            if s.holding is None:
+                return "FAILED_EMPTY"
+            bin_name = "A" if action == "PLACE_BIN_A" else "B"
+            obj = s.objects[s.holding]
+            # If this object was blocking something, reveal that target now.
+            if obj.blocking and obj.blocking in s.objects:
+                s.objects[obj.blocking].reachable = True
+                obj.blocking = None
+            obj.in_bin = bin_name
+            obj.is_held = False
+            obj.reachable = False
+            s.holding = None
+            s.gripper_open = True
+            return "SUCCESS"
+
+        elif action == "CLEAR_BLOCKER":
+            for obj in s.objects.values():
+                if obj.blocking is not None and obj.reachable:
+                    blocked_name = obj.blocking
+                    obj.blocking = None
+                    obj.pos = obj.pos + np.array([0.28, 0.1, 0])
+                    if blocked_name in s.objects:
+                        s.objects[blocked_name].reachable = True
+                    return "SUCCESS"
+            return "FAILED_INVALID"
+
+        return "FAILED_INVALID"
+
+    def get_state(self) -> SimState:
+        return self._state
+
+    def get_camera_obs(self) -> Optional[dict]:
+        """
+        Return camera observations + vision-extracted symbolic state.
+
+        Stub mode:   returns None (symbolic state comes from sim config directly)
+        Real mode:   returns RGB images + runs vision.py to extract object positions
+
+        The planning layer above never needs to know which path ran —
+        it always receives the same symbolic SimState either way.
+        """
+        if self.use_stub:
+            return None  # stub: symbolic state already in self._state, no camera needed
+
+        if self._env is not None:
+            obs = self._env._get_observations()
+            rgb_front = obs.get("frontview_image")
+            rgb_agent = obs.get("agentview_image")
+
+            # Run vision pipeline to get symbolic state from images
+            if rgb_front is not None and self._current_cfg is not None:
+                from .vision import sim_vision
+                vision_result = sim_vision(rgb_front)
+                # Merge detected positions back into symbolic state
+                # (perception layer updates what was set from physics)
+                for det in vision_result.detected_objects:
+                    name = det["name"]
+                    if name in self._state.objects:
+                        self._state.objects[name].pos = np.array([det["x"], det["y"], det["z"]])
+
+            return {"frontview": rgb_front, "agentview": rgb_agent}
+
+
+# ── Re-exports ─────────────────────────────────────────────────────────
+
+__all__ = ["SimWrapper", "SimState", "ObjectState"]