Upload folder using huggingface_hub

models.py

@@ -1,27 +1,61 @@
 from __future__ import annotations
 
+import json
 from typing import Any
 
-from pydantic import BaseModel, Field
-
-
-class AgenticTrafficAction(BaseModel):
-    district_actions: dict[str, Any] = Field(default_factory=dict)
-
-
-class AgenticTrafficObservation(BaseModel):
+from openenv.core.env_server import Action, Observation, State
+from pydantic import Field, field_validator
+
+
+class AgenticTrafficAction(Action):
+    use_llm: bool = Field(
+        default=False,
+        description=(
+            "When true, use the bundled district LLM adapter to generate district_actions "
+            "for districts not explicitly provided."
+        ),
+    )
+    district_actions: dict[str, Any] = Field(
+        default_factory=dict,
+        description=(
+            "JSON object keyed by district_id. Use {} for a no-op step, or provide "
+            'entries like {"d_00":{"strategy":"hold","phase_bias":"NS","duration_steps":10}}.'
+        ),
+    )
+    llm_max_new_tokens: int = Field(
+        default=128,
+        ge=16,
+        le=512,
+        description="Maximum new tokens to generate per district when use_llm=true.",
+    )
+
+    @field_validator("district_actions", mode="before")
+    @classmethod
+    def parse_district_actions(cls, value: Any) -> dict[str, Any]:
+        if value is None or value == "":
+            return {}
+        if isinstance(value, str):
+            parsed = json.loads(value)
+            if not isinstance(parsed, dict):
+                raise ValueError("district_actions must decode to a JSON object.")
+            return parsed
+        if isinstance(value, dict):
+            return value
+        raise ValueError("district_actions must be a dict or JSON object string.")
+
+
+class AgenticTrafficObservation(Observation):
     city_id: str | None = None
     scenario_name: str | None = None
     decision_step: int = 0
     sim_time: int = 0
     district_summaries: dict[str, Any] = Field(default_factory=dict)
-    done: bool = False
-    reward: float = 0.0
 
 
-class AgenticTrafficState(BaseModel):
+class AgenticTrafficState(State):
     scenario: dict[str, Any] | None = None
     controller: dict[str, Any] = Field(default_factory=dict)
     district_decision_interval: int = 0
     district_summaries: dict[str, Any] = Field(default_factory=dict)
+    llm: dict[str, Any] = Field(default_factory=dict)
     last_info: dict[str, Any] = Field(default_factory=dict)

server/environment.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
 import os
+import uuid
 from pathlib import Path
+from typing import Any
 
+from district_llm.inference import DistrictLLMInference
+from district_llm.schema import DistrictAction
 from models import (
     AgenticTrafficAction,
     AgenticTrafficObservation,
@@ -15,6 +19,11 @@ from openenv_app.openenv_wrapper import OpenEnvTrafficWrapper
 REPO_ROOT = Path(__file__).resolve().parents[1]
 DATA_DIR = Path(os.environ.get("DATA_DIR", "") or (REPO_ROOT / "data" / "generated"))
 SPLITS_DIR = Path(os.environ.get("SPLITS_DIR", "") or (REPO_ROOT / "data" / "splits"))
+DISTRICT_LLM_ADAPTER_PATH = Path(
+    os.environ.get("DISTRICT_LLM_ADAPTER_PATH", "")
+    or (REPO_ROOT / "artifacts" / "district_llm_adapter_v3" / "main_run" / "adapter")
+)
+DISTRICT_LLM_DEVICE = os.environ.get("DISTRICT_LLM_DEVICE")
 
 
 class AgenticTrafficEnvironment(
@@ -23,23 +32,50 @@ class AgenticTrafficEnvironment(
     """Minimal OpenEnv-compatible wrapper around the existing district controller stack."""
 
     def __init__(self) -> None:
+        super().__init__()
         self.wrapper = OpenEnvTrafficWrapper(
             generated_root=DATA_DIR,
             splits_root=SPLITS_DIR,
         )
         self._state = AgenticTrafficState()
+        self._llm_inference: DistrictLLMInference | None = None
+        self._llm_load_attempted = False
+        self._llm_error: str | None = None
 
-    def reset(self) -> AgenticTrafficObservation:
-        payload = self.wrapper.reset(seed=None)
+    def reset(
+        self,
+        seed: int | None = None,
+        episode_id: str | None = None,
+        **kwargs: Any,
+    ) -> AgenticTrafficObservation:
+        payload = self.wrapper.reset(
+            seed=seed,
+            city_id=kwargs.get("city_id"),
+            scenario_name=kwargs.get("scenario_name"),
+        )
+        self._state.episode_id = episode_id or str(uuid.uuid4())
+        self._state.step_count = 0
         self._sync_state()
-        return AgenticTrafficObservation.model_validate(payload["observation"])
+        observation = AgenticTrafficObservation.model_validate(payload["observation"])
+        observation.reward = None
+        observation.done = False
+        observation.metadata["llm"] = self._llm_status()
+        return observation
 
-    def step(self, action: AgenticTrafficAction) -> AgenticTrafficObservation:
-        payload = self.wrapper.step(action=action.model_dump())
+    def step(
+        self,
+        action: AgenticTrafficAction,
+        timeout_s: float | None = None,
+        **kwargs: Any,
+    ) -> AgenticTrafficObservation:
+        del timeout_s, kwargs
+        payload = self.wrapper.step(action=self._build_step_payload(action))
+        self._state.step_count += 1
         self._sync_state()
         observation = AgenticTrafficObservation.model_validate(payload["observation"])
         observation.done = bool(payload.get("done", False))
         observation.reward = float(payload.get("reward", 0.0))
+        observation.metadata["llm"] = self._llm_status()
         return observation
 
     @property
@@ -47,6 +83,86 @@ class AgenticTrafficEnvironment(
         self._sync_state()
         return self._state
 
+    def _build_step_payload(self, action: AgenticTrafficAction) -> dict[str, Any]:
+        district_actions = dict(action.district_actions)
+        llm_generated_actions: dict[str, Any] = {}
+
+        if action.use_llm:
+            llm_generated_actions = self._generate_llm_actions(
+                existing_actions=district_actions,
+                max_new_tokens=action.llm_max_new_tokens,
+            )
+            for district_id, directive in llm_generated_actions.items():
+                district_actions.setdefault(district_id, directive)
+
+        payload = {"district_actions": district_actions}
+        payload["metadata"] = {
+            "use_llm": bool(action.use_llm),
+            "llm_generated_districts": sorted(llm_generated_actions),
+            "llm": self._llm_status(),
+        }
+        return payload
+
+    def _generate_llm_actions(
+        self,
+        existing_actions: dict[str, Any],
+        max_new_tokens: int,
+    ) -> dict[str, Any]:
+        if not self.wrapper.last_summaries:
+            return {}
+
+        inference = self._get_llm_inference()
+        if inference is None:
+            return {}
+
+        generated_actions: dict[str, Any] = {}
+        for district_id, summary in self.wrapper.last_summaries.items():
+            if district_id in existing_actions:
+                continue
+            result = inference.predict_with_result(summary=summary, max_new_tokens=max_new_tokens)
+            generated_actions[district_id] = result.action.to_dict()
+        return generated_actions
+
+    def _get_llm_inference(self) -> DistrictLLMInference | None:
+        if self._llm_inference is not None:
+            return self._llm_inference
+        if self._llm_load_attempted:
+            return None
+
+        self._llm_load_attempted = True
+        if not DISTRICT_LLM_ADAPTER_PATH.exists():
+            self._llm_error = f"Adapter not found at {DISTRICT_LLM_ADAPTER_PATH}"
+            return None
+
+        try:
+            self._llm_inference = DistrictLLMInference(
+                model_name_or_path=str(DISTRICT_LLM_ADAPTER_PATH),
+                device=DISTRICT_LLM_DEVICE,
+                fallback_action=DistrictAction.default_hold(
+                    duration_steps=self.wrapper.district_decision_interval
+                ),
+            )
+        except Exception as exc:
+            self._llm_error = f"{type(exc).__name__}: {exc}"
+            self._llm_inference = None
+        return self._llm_inference
+
+    def _llm_status(self) -> dict[str, Any]:
+        return {
+            "adapter_path": str(DISTRICT_LLM_ADAPTER_PATH),
+            "adapter_present": DISTRICT_LLM_ADAPTER_PATH.exists(),
+            "loaded": self._llm_inference is not None,
+            "load_attempted": self._llm_load_attempted,
+            "error": self._llm_error,
+        }
+
     def _sync_state(self) -> None:
         payload = self.wrapper.state()["state"]
-        self._state = AgenticTrafficState.model_validate(payload)
+        self._state = AgenticTrafficState.model_validate(
+            {
+                **payload,
+                "episode_id": self._state.episode_id,
+                "step_count": self._state.step_count,
+                "llm": self._llm_status(),
+            }
+        )