Upload server/app.py with huggingface_hub

server/app.py

@@ -0,0 +1,340 @@
+"""
+RoboReplan server — OpenEnv HTTP protocol + metrics endpoint.
+"""
+import os
+import re
+from pathlib import Path
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse, RedirectResponse
+from pydantic import BaseModel
+from openenv.core.env_server import create_fastapi_app
+
+from .openenv_env import RoboReplanEnv, RoboAction, RoboObservation, RoboState
+from .models import Action as EnvAction
+
+difficulty = os.environ.get("DIFFICULTY", "easy")
+
+# Shared env instance (metrics persist across requests)
+_env_instance = RoboReplanEnv(difficulty=difficulty)
+
+app = create_fastapi_app(
+    env=lambda: _env_instance,
+    action_cls=RoboAction,
+    observation_cls=RoboObservation,
+)
+
+app.add_middleware(CORSMiddleware, allow_origins=["*"],
+                   allow_methods=["*"], allow_headers=["*"])
+
+_VIZ_HTML = (Path(__file__).parent.parent / "viz_standalone.html").read_text()
+
+
+@app.get("/")
+def root():
+    return RedirectResponse(url="/viz")
+
+
+@app.get("/viz", response_class=HTMLResponse)
+def viz():
+    return _VIZ_HTML
+
+
+@app.get("/metrics")
+def metrics():
+    """Live training metrics: success rate, reward curve, failure breakdown, oracle agreement."""
+    return _env_instance.metrics
+
+
+# ── Demo endpoints — judges can interact live ──────────────────────────
+
+_demo_env = None
+_policy_pipe = None
+_POLICY_MODEL = os.environ.get("DEMO_POLICY_MODEL", "jshah13/robo-replan-grpo")
+_VALID_ACTIONS = [a.value for a in EnvAction]
+
+
+def _extract_action(text: str) -> str:
+    clean = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip().upper()
+    normalized = re.sub(r"[^A-Z_ ]+", " ", clean)
+    normalized = re.sub(r"\s+", " ", normalized).strip()
+    for action in sorted(_VALID_ACTIONS, key=len, reverse=True):
+        if action in clean:
+            return action
+    spaced_map = {a.replace("_", " "): a for a in _VALID_ACTIONS}
+    for spaced, action in spaced_map.items():
+        if spaced in normalized:
+            return action
+    return "SCAN_SCENE"
+
+
+def _parse_valid_actions_from_prompt(prompt: str) -> list[str]:
+    m = re.search(r"Valid now:\s*(.*)", prompt, flags=re.IGNORECASE)
+    if not m:
+        return []
+    raw = m.group(1).strip()
+    if raw.lower() == "any":
+        return []
+    items = [x.strip().upper() for x in raw.split(",") if x.strip()]
+    return [a for a in items if a in _VALID_ACTIONS]
+
+
+def _fallback_action(valid: list[str]) -> str:
+    if not valid:
+        return "SCAN_SCENE"
+    priority = [
+        "PLACE_BIN_A", "PLACE_BIN_B",
+        "PICK",
+        "CLEAR_BLOCKER",
+        "MOVE_TO_RED", "MOVE_TO_BLUE", "MOVE_TO_GREEN", "MOVE_TO_YELLOW", "MOVE_TO_PURPLE",
+        "MOVE_NORTH", "MOVE_SOUTH", "MOVE_EAST", "MOVE_WEST",
+        "ROTATE_LEFT", "ROTATE_RIGHT",
+        "SCAN_SCENE",
+    ]
+    for p in priority:
+        if p in valid:
+            return p
+    return valid[0]
+
+
+def _prompt_line(prompt: str, key: str) -> str:
+    m = re.search(rf"{re.escape(key)}:\s*(.*)", prompt, flags=re.IGNORECASE)
+    return m.group(1).strip() if m else ""
+
+
+def _smart_fallback_action(valid: list[str], prompt: str) -> str:
+    """
+    Use lightweight state cues from prompt to avoid repetitive bad actions.
+    """
+    if not valid:
+        return "SCAN_SCENE"
+    valid_set = set(valid)
+    last_line = _prompt_line(prompt, "Last")
+    holding = _prompt_line(prompt, "Holding").lower()
+
+    last_action, last_result = "", ""
+    m = re.match(r"\s*([A-Z_]+)\s*->\s*([A-Z_]+)", last_line.upper())
+    if m:
+        last_action, last_result = m.group(1), m.group(2)
+
+    # If holding something, prioritize placing over anything else.
+    if holding and holding not in ("nothing", "none", "null"):
+        if "PLACE_BIN_A" in valid_set:
+            return "PLACE_BIN_A"
+        if "PLACE_BIN_B" in valid_set:
+            return "PLACE_BIN_B"
+
+    # If we just moved to a target successfully, then pick.
+    if last_action.startswith("MOVE_TO_") and last_result == "SUCCESS" and "PICK" in valid_set:
+        return "PICK"
+
+    # If pick just failed/was invalid, move or clear first instead of repeating pick.
+    if last_action == "PICK" and last_result.startswith("FAILED"):
+        for a in valid:
+            if a.startswith("MOVE_TO_"):
+                return a
+        if "CLEAR_BLOCKER" in valid_set:
+            return "CLEAR_BLOCKER"
+
+    # In top-down mode, blind PICK tends to loop; prefer moving to a target first.
+    for a in valid:
+        if a.startswith("MOVE_TO_"):
+            return a
+
+    return _fallback_action(valid)
+
+
+def _get_policy_pipe():
+    global _policy_pipe
+    if _policy_pipe is None:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+        tokenizer = AutoTokenizer.from_pretrained(_POLICY_MODEL)
+        tokenizer.padding_side = "left"
+        tokenizer.pad_token = tokenizer.eos_token
+
+        has_gpu = torch.cuda.is_available()
+        if has_gpu:
+            model = AutoModelForCausalLM.from_pretrained(
+                _POLICY_MODEL, torch_dtype=torch.float16, device_map="auto",
+            )
+            pipe_kwargs = {"device_map": "auto"}
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                _POLICY_MODEL, torch_dtype=torch.float32,
+            )
+            pipe_kwargs = {"device": "cpu"}
+
+        model.generation_config.pad_token_id = tokenizer.pad_token_id
+        model.generation_config.max_length = None
+        _policy_pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, **pipe_kwargs)
+    return _policy_pipe
+
+
+class PolicyActionRequest(BaseModel):
+    prompt: str
+    valid_actions: list[str] = []
+
+
+def _format_demo_step_response(step_out):
+    """
+    Compat for env implementations that return either:
+    - StepResult(observation, reward, done, info), or
+    - Observation(done, reward, ...)
+    """
+    if hasattr(step_out, "observation"):
+        obs = step_out.observation
+        return {
+            "observation": obs.model_dump(),
+            "reward": float(getattr(step_out, "reward", 0.0) or 0.0),
+            "done": bool(getattr(step_out, "done", False)),
+            "info": getattr(step_out, "info", {}) or {},
+        }
+    obs = step_out
+    return {
+        "observation": obs.model_dump(),
+        "reward": float(getattr(obs, "reward", 0.0) or 0.0),
+        "done": bool(getattr(obs, "done", False)),
+        "info": {},
+    }
+
+
+@app.post("/demo/reset")
+def demo_reset(difficulty: str = "easy", scenario_pack: str = "default"):
+    """Start a fresh demo episode. scenario_pack: default | pharmacy | warehouse | lab"""
+    global _demo_env
+    _demo_env = RoboReplanEnv(difficulty=difficulty)
+    # Apply scenario pack by patching the internal env config before reset
+    if scenario_pack != "default":
+        _demo_env._env.cfg.task.scenario_pack = scenario_pack
+    obs = _demo_env.reset()
+    return {"observation": obs.model_dump(), "done": False, "reward": 0.0}
+
+
+@app.post("/demo/step")
+def demo_step(action: str):
+    """Take one step in the demo episode."""
+    global _demo_env
+    if _demo_env is None:
+        _demo_env = RoboReplanEnv(difficulty="easy")
+        _demo_env.reset()
+    result = _demo_env.step(RoboAction(action=action))
+    return _format_demo_step_response(result)
+
+
+def _oracle_reasoning(env, action: str) -> str:
+    """Generate a human-readable <think> narration for the oracle's chosen action."""
+    try:
+        inner = env._env
+        state = inner.sim.get_state()
+        holding = state.holding
+        placements = inner._required_placements
+        failures = inner._known_failures
+        constraints = inner._active_constraints
+        completed = set(inner._completed_subgoals)
+        instruction = inner._instruction
+
+        # Build context strings
+        blocked_objs = [n for n, o in state.objects.items() if not o.reachable and o.in_bin is None]
+        remaining = [n for n, b in placements.items()
+                     if f"placed_{n}_in_bin_{b}" not in completed and
+                     not (state.objects.get(n) and state.objects[n].in_bin == b)]
+        constraint_str = f" Constraint: {', '.join(constraints)}." if constraints else ""
+        deadline_status = inner._deadline_status()
+        deadline_str = f" Deadlines active: {deadline_status}." if deadline_status else ""
+
+        if action == "SCAN_SCENE":
+            return (f"I need to survey the scene to reveal hidden object traits "
+                    f"before planning.{constraint_str}")
+        if action == "CLEAR_BLOCKER":
+            for n, o in state.objects.items():
+                if o.blocking and o.reachable:
+                    return (f"Plan: CLEAR_BLOCKER → MOVE_TO_{o.blocking.replace('_block','').upper()} → PICK → PLACE. "
+                            f"{n} is blocking {o.blocking}. Clearing it first.{constraint_str}")
+        if holding and action.startswith("PLACE_BIN_"):
+            bin_name = action.split("_")[-1]
+            target_bin = placements.get(holding, "?")
+            correct = target_bin == bin_name
+            return (f"I am holding {holding}. Target bin is {target_bin}. "
+                    f"{'Placing correctly' if correct else 'Placing in available bin'}.{constraint_str}")
+        if action.startswith("MOVE_TO_"):
+            color = action.replace("MOVE_TO_", "").lower()
+            target = f"{color}_block"
+            if failures:
+                return (f"Previous failure: {failures[-1]}. Re-navigating to {target} to retry.{constraint_str}")
+            return (f"Plan: MOVE_TO_{color.upper()} → PICK → PLACE_BIN_{placements.get(target,'?')}. "
+                    f"Moving to {target} now.{constraint_str}{deadline_str}")
+        if action == "PICK":
+            if failures and any("FAILED_SLIP" in f for f in failures):
+                return (f"Previous grasp slip detected ({failures[-1]}). "
+                        f"Retrying PICK — repositioning and attempting again.")
+            return (f"Gripper is adjacent to target. Attempting to grasp.{constraint_str}")
+        if action in ("MOVE_NORTH", "MOVE_SOUTH", "MOVE_EAST", "MOVE_WEST"):
+            if remaining:
+                target = remaining[0]
+                return (f"Navigating toward {target} "
+                        f"(target bin: {placements.get(target,'?')}).{constraint_str}{deadline_str}")
+        return f"Executing {action}.{constraint_str}"
+    except Exception:
+        return f"Executing {action}."
+
+
+@app.get("/demo/oracle")
+def demo_oracle():
+    """Step using the oracle policy — shows optimal behavior for demo."""
+    global _demo_env
+    if _demo_env is None:
+        _demo_env = RoboReplanEnv(difficulty="easy")
+        _demo_env.reset()
+    oracle = _demo_env._env._oracle_action() or "SCAN_SCENE"
+    reasoning = _oracle_reasoning(_demo_env, oracle)
+    result = _demo_env.step(RoboAction(action=oracle))
+    payload = _format_demo_step_response(result)
+    payload["action_taken"] = oracle
+    payload["reasoning"] = reasoning
+    return payload
+
+
+@app.post("/demo/policy_action")
+def demo_policy_action(req: PolicyActionRequest):
+    """
+    Returns one model-predicted action for a given prompt.
+    The visualization can use this to drive the environment step-by-step.
+    """
+    try:
+        pipe = _get_policy_pipe()
+        out = pipe(
+            req.prompt,
+            return_full_text=False,
+            max_new_tokens=128,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.05,
+        )[0]["generated_text"]
+        raw_action = _extract_action(out)
+        action = raw_action
+        valid = [a for a in req.valid_actions if a in _VALID_ACTIONS] or _parse_valid_actions_from_prompt(req.prompt)
+        if valid and action not in valid:
+            action = _smart_fallback_action(valid, req.prompt)
+        # Avoid no-op scan loops when other valid actions exist.
+        if valid and action == "SCAN_SCENE" and any(v != "SCAN_SCENE" for v in valid):
+            action = _smart_fallback_action([v for v in valid if v != "SCAN_SCENE"], req.prompt)
+        # Extract <think>...</think> reasoning separately for display and env reward
+        import re as _re
+        _m = _re.search(r'<think>(.*?)</think>', out, _re.DOTALL)
+        reasoning = _m.group(1).strip() if _m else ""
+        return {
+            "action": action,
+            "reasoning": reasoning,
+            "raw_output": out,
+            "raw_action": raw_action,
+            "valid_actions_used": valid,
+        }
+    except Exception as exc:
+        # Fail soft so the UI can still run with manual/scripted controls.
+        valid = [a for a in req.valid_actions if a in _VALID_ACTIONS] or _parse_valid_actions_from_prompt(req.prompt)
+        action = _smart_fallback_action([v for v in valid if v != "SCAN_SCENE"], req.prompt) if valid else "SCAN_SCENE"
+        return {"action": action, "error": str(exc), "valid_actions_used": valid}