Spaces:

1akshith
/

interior-env

Sleeping

Interior Env

Added LOCAL_IMAGE_NAME to environment vars to perfectly satisfy the pre-submission static validator

665afbb 7 months ago

7.2 kB

"""
inference.py — OpenEnv-compliant inference entry point for the
Interior Design Spatial Engine (Furniture Tetris + Human Flow).
STDOUT FORMAT (strictly followed):
    [START] task=<task> env=<benchmark> model=<model>
    [STEP]  step=<n> action=<str> reward=<0.00> done=<true|false> error=<msg|null>
    [END]   success=<true|false> steps=<n> score=<0.000> rewards=<r1,r2,...>
"""

from __future__ import annotations
import sys
import os
import json
import argparse
from typing import List, Optional

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from env import InteriorEnv
from baseline import BaselineAgent
from models import Action
from tasks import TASKS

# ── Environment variables ─────────────────────────────────────────────────────
API_BASE_URL     = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME       = os.getenv("MODEL_NAME",   "Qwen/Qwen2.5-72B-Instruct")
HF_TOKEN         = os.getenv("HF_TOKEN")
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
BENCHMARK    = "interior-env"

# ── OpenAI client (optional LLM agent) ───────────────────────────────────────
try:
    from openai import OpenAI
    _llm_client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN) if HF_TOKEN else None
except Exception:
    _llm_client = None

# ── Strict stdout loggers ─────────────────────────────────────────────────────
def log_start(task: str, model: str) -> None:
    print(f"[START] task={task} env={BENCHMARK} model={model}", flush=True)

def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
    err = error if error else "null"
    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={err}", flush=True)

def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    rew_str = ",".join(f"{r:.2f}" for r in rewards)
    print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rew_str}", flush=True)

# ── LLM Agent ────────────────────────────────────────────────────────────────
SYSTEM_PROMPT = """You are a spatial AI agent designing interior layouts on a 2D grid.
You place furniture ensuring strict geometrical, human flow, and physics constraints.
- Light rays shoot purely orthogonally from windows. Tall items block them (cast shadows).
- Windows/doors emit cold (-2). Radiators/TVs emit radial heat. 
Output ONLY valid JSON, no explanation:
{"action_type": "place_item", "item": "<item_id>", "x": <int>, "y": <int>, "rotation": <0|90>}"""

def llm_pick_action(obs, step: int) -> Optional[Action]:
    if not _llm_client:
        return None
    try:
        obs_summary = {
            "room_size": obs.room_size,
            "door_position": obs.door_position,
            "windows": obs.windows,
            "unplaced_items": [{"id": u.id, "width": u.width, "height": u.height, "rules": u.rules}
                               for u in obs.unplaced_items],
            "placed_items": {k: {"x": v.x, "y": v.y, "rotation": v.rotation}
                            for k, v in obs.placed_items.items()},
            "feedback": obs.feedback_logs[-3:] if obs.feedback_logs else [],
        }
        completion = _llm_client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user",   "content": f"Step {step}. Current state:\n{json.dumps(obs_summary, indent=2)}\n\nWhat is your next action?"},
            ],
            response_format={"type": "json_object"},
            temperature=0.2,
            max_tokens=100,
        )
        data = json.loads(completion.choices[0].message.content)
        return Action(
            action_type=data.get("action_type", "no_op"),
            item=data.get("item", ""),
            x=int(data.get("x", 0)),
            y=int(data.get("y", 0)),
            rotation=int(data.get("rotation", 0)),
        )
    except Exception as e:
        return None

# ── Action → compact string ─────────────────────────────────────────────────
def action_str(action: Action) -> str:
    if action.action_type == "place_item":
        return f"place_item({action.item},{action.x},{action.y},{action.rotation})"
    if action.action_type == "remove_item":
        return f"remove_item({action.item})"
    return "no_op()"

# ── Single episode runner ─────────────────────────────────────────────────────
def run_task(task_id: str) -> None:
    if task_id not in TASKS:
        print(f"[ERROR] Unknown task: {task_id}", flush=True)
        return

    config   = TASKS[task_id]
    env      = InteriorEnv()
    baseline = BaselineAgent()

    obs  = env.reset(config, task_id)
    done = False
    step = 0
    rewards: List[float] = []
    success = False
    score   = 0.0

    log_start(task=task_id, model=MODEL_NAME)

    try:
        while not done:
            step += 1

            # Try LLM first, fall back to deterministic baseline
            action = llm_pick_action(obs, step)
            if action is None:
                action = baseline.pick_action(obs, task_id)

            obs, reward, done, _ = env.step(action)
            rewards.append(reward)

            error = None
            if obs.feedback_logs:
                last = obs.feedback_logs[-1]
                if any(tag in last for tag in ("ERROR", "COLLISION", "BOUNDS", "FLOW")):
                    error = last[:80].replace(" ", "_")

            log_step(step=step, action=action_str(action),
                     reward=reward, done=done, error=error)

        result  = env.get_episode_result()
        success = result.success
        score   = result.score

    except Exception as exc:
        log_step(step=step + 1, action="error()", reward=0.0, done=True, error=str(exc)[:80])

    finally:
        log_end(success=success, steps=step, score=score, rewards=rewards)

# ── Entry point ───────────────────────────────────────────────────────────────
def main() -> None:
    parser = argparse.ArgumentParser(description="Interior Env OpenEnv inference")
    parser.add_argument("--task", choices=list(TASKS.keys()),
                        help="Run a single task. Omit to run all levels.")
    args = parser.parse_args()

    tasks_to_run = [args.task] if args.task else list(TASKS.keys())
    for tid in tasks_to_run:
        run_task(tid)

if __name__ == "__main__":
    main()

	"""
	inference.py — OpenEnv-compliant inference entry point for the
	Interior Design Spatial Engine (Furniture Tetris + Human Flow).

	STDOUT FORMAT (strictly followed):
	[START] task=<task> env=<benchmark> model=<model>
	[STEP] step=<n> action=<str> reward=<0.00> done=<true\|false> error=<msg\|null>
	[END] success=<true\|false> steps=<n> score=<0.000> rewards=<r1,r2,...>
	"""

	from __future__ import annotations
	import sys
	import os
	import json
	import argparse
	from typing import List, Optional

	sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

	from env import InteriorEnv
	from baseline import BaselineAgent
	from models import Action
	from tasks import TASKS

	# ── Environment variables ─────────────────────────────────────────────────────
	API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
	MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
	HF_TOKEN = os.getenv("HF_TOKEN")
	LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
	BENCHMARK = "interior-env"

	# ── OpenAI client (optional LLM agent) ───────────────────────────────────────
	try:
	from openai import OpenAI
	_llm_client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN) if HF_TOKEN else None
	except Exception:
	_llm_client = None

	# ── Strict stdout loggers ─────────────────────────────────────────────────────
	def log_start(task: str, model: str) -> None:
	print(f"[START] task={task} env={BENCHMARK} model={model}", flush=True)

	def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
	err = error if error else "null"
	print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={err}", flush=True)

	def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
	rew_str = ",".join(f"{r:.2f}" for r in rewards)
	print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rew_str}", flush=True)

	# ── LLM Agent ────────────────────────────────────────────────────────────────
	SYSTEM_PROMPT = """You are a spatial AI agent designing interior layouts on a 2D grid.
	You place furniture ensuring strict geometrical, human flow, and physics constraints.
	- Light rays shoot purely orthogonally from windows. Tall items block them (cast shadows).
	- Windows/doors emit cold (-2). Radiators/TVs emit radial heat.
	Output ONLY valid JSON, no explanation:
	{"action_type": "place_item", "item": "<item_id>", "x": <int>, "y": <int>, "rotation": <0\|90>}"""

	def llm_pick_action(obs, step: int) -> Optional[Action]:
	if not _llm_client:
	return None
	try:
	obs_summary = {
	"room_size": obs.room_size,
	"door_position": obs.door_position,
	"windows": obs.windows,
	"unplaced_items": [{"id": u.id, "width": u.width, "height": u.height, "rules": u.rules}
	for u in obs.unplaced_items],
	"placed_items": {k: {"x": v.x, "y": v.y, "rotation": v.rotation}
	for k, v in obs.placed_items.items()},
	"feedback": obs.feedback_logs[-3:] if obs.feedback_logs else [],
	}
	completion = _llm_client.chat.completions.create(
	model=MODEL_NAME,
	messages=[
	{"role": "system", "content": SYSTEM_PROMPT},
	{"role": "user", "content": f"Step {step}. Current state:\n{json.dumps(obs_summary, indent=2)}\n\nWhat is your next action?"},
	],
	response_format={"type": "json_object"},
	temperature=0.2,
	max_tokens=100,
	)
	data = json.loads(completion.choices[0].message.content)
	return Action(
	action_type=data.get("action_type", "no_op"),
	item=data.get("item", ""),
	x=int(data.get("x", 0)),
	y=int(data.get("y", 0)),
	rotation=int(data.get("rotation", 0)),
	)
	except Exception as e:
	return None

	# ── Action → compact string ─────────────────────────────────────────────────
	def action_str(action: Action) -> str:
	if action.action_type == "place_item":
	return f"place_item({action.item},{action.x},{action.y},{action.rotation})"
	if action.action_type == "remove_item":
	return f"remove_item({action.item})"
	return "no_op()"

	# ── Single episode runner ─────────────────────────────────────────────────────
	def run_task(task_id: str) -> None:
	if task_id not in TASKS:
	print(f"[ERROR] Unknown task: {task_id}", flush=True)
	return

	config = TASKS[task_id]
	env = InteriorEnv()
	baseline = BaselineAgent()

	obs = env.reset(config, task_id)
	done = False
	step = 0
	rewards: List[float] = []
	success = False
	score = 0.0

	log_start(task=task_id, model=MODEL_NAME)

	try:
	while not done:
	step += 1

	# Try LLM first, fall back to deterministic baseline
	action = llm_pick_action(obs, step)
	if action is None:
	action = baseline.pick_action(obs, task_id)

	obs, reward, done, _ = env.step(action)
	rewards.append(reward)

	error = None
	if obs.feedback_logs:
	last = obs.feedback_logs[-1]
	if any(tag in last for tag in ("ERROR", "COLLISION", "BOUNDS", "FLOW")):
	error = last[:80].replace(" ", "_")

	log_step(step=step, action=action_str(action),
	reward=reward, done=done, error=error)

	result = env.get_episode_result()
	success = result.success
	score = result.score

	except Exception as exc:
	log_step(step=step + 1, action="error()", reward=0.0, done=True, error=str(exc)[:80])

	finally:
	log_end(success=success, steps=step, score=score, rewards=rewards)

	# ── Entry point ───────────────────────────────────────────────────────────────
	def main() -> None:
	parser = argparse.ArgumentParser(description="Interior Env OpenEnv inference")
	parser.add_argument("--task", choices=list(TASKS.keys()),
	help="Run a single task. Omit to run all levels.")
	args = parser.parse_args()

	tasks_to_run = [args.task] if args.task else list(TASKS.keys())
	for tid in tasks_to_run:
	run_task(tid)

	if __name__ == "__main__":
	main()