Upload folder using huggingface_hub

src/__init__.py

@@ -0,0 +1,10 @@
+import logging
+
+
+class EndpointFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        return record.getMessage().find("/connect_four/bot_progress") == -1
+
+
+# Filter out /endpoint
+logging.getLogger("uvicorn.access").addFilter(EndpointFilter())

src/app.py

@@ -1,49 +1,59 @@
-from typing import Annotated, Any, Generator
+from collections.abc import Generator
 from pathlib import Path
-from gymnasium.wrappers.record_video import RecordVideo
-import numpy as np
-from litrl.env.make import make
-from litrl.common.agent import RandomAgent
-from litrl.env.typing import SingleAgentId
+from typing import Annotated, Any
+
 from fastapi import Depends, FastAPI, Request, status
 from fastapi.exceptions import RequestValidationError
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
-from litrl.env.connect_four import Board
+from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse
+from gymnasium.wrappers.record_video import RecordVideo
 from loguru import logger
-from fastapi.responses import StreamingResponse, RedirectResponse
-from src.app_state import AppState
-from src.typing import CpuConfig
 from pydantic import BaseModel
+from src.app_state import AppState
 from src.huggingface.huggingface_client import HuggingFaceClient
+from src.typing import CpuConfig
+
+from litrl.algo.mcts.agent import MCTSAgent
+from litrl.common.agent import RandomAgent
 from litrl.env.connect_four import ConnectFour
+from litrl.env.make import make
+from litrl.env.typing import GymId
+
 
 def stream_mp4(mp4_path: Path) -> StreamingResponse:
-    def iter_file()-> Generator[bytes, Any, None]:  
+    def iter_file() -> Generator[bytes, Any, None]:
         with mp4_path.open(mode="rb") as env_file:
             yield from env_file
 
     return StreamingResponse(content=iter_file(), media_type="video/mp4")
 
+
 ObservationType = list[list[list[int]]]
+
+
 class GridResponseType(BaseModel):
     grid: ObservationType
     done: bool
 
-def step(env: ConnectFour, action: int)->GridResponseType:
+
+def step(env: ConnectFour, action: int) -> GridResponseType:
     env.step(action)
     return observe(env)
 
-def observe(env: ConnectFour)->GridResponseType:
+
+def observe(env: ConnectFour) -> GridResponseType:
     obs = env.observe("player_1")
-    done = env.terminations[env.agent_selection] or env.truncations[env.agent_selection]
-    return {"grid": obs['observation'].tolist(), "done": done}
+    return GridResponseType(
+        grid=obs["observation"].tolist(),
+        done=bool(env.terminations[env.agent_selection] or env.truncations[env.agent_selection]),  # TODO why needed?
+    )
+
 
-def create_app() -> FastAPI:
+def create_app() -> FastAPI:  # noqa: C901 # TODO move to routes
     app = FastAPI()
 
-    @app.get('/')
-    async def to_docs():
+    @app.get("/")
+    async def redirect_to_docs() -> RedirectResponse:
         return RedirectResponse("/docs")
 
     @app.post(path="/connect_four/play", response_model=GridResponseType)
@@ -52,13 +62,13 @@ def create_app() -> FastAPI:
         app_state: Annotated[AppState, Depends(dependency=AppState)],
     ) -> GridResponseType:
         return step(app_state.env, action)
-    
+
     @app.get(path="/connect_four/observe", response_model=GridResponseType)
     def endpoint_observe(
         app_state: Annotated[AppState, Depends(dependency=AppState)],
     ) -> GridResponseType:
         return observe(app_state.env)
-    
+
     @app.post(path="/connect_four/bot_play", response_model=GridResponseType)
     def endpoint_bot_play(
         cpu_config: CpuConfig,
@@ -68,6 +78,18 @@ def create_app() -> FastAPI:
         action = app_state.get_action()
         return step(app_state.env, action)
 
+    @app.get(path="/connect_four/bot_progress", response_model=float)
+    def endpoint_bot_progress(
+        app_state: Annotated[AppState, Depends(dependency=AppState)],
+    ) -> float:
+        if isinstance(app_state.agent, MCTSAgent):
+            if app_state.cpu_config.simulations is None:
+                raise ValueError
+            return float(
+                app_state.agent.mcts._root.visits / app_state.cpu_config.simulations,  # noqa: SLF001
+            )  # TODO why needed?
+        return 1.0
+
     @app.get(path="/connect_four/reset", response_model=GridResponseType)
     def endpoint_reset(
         app_state: Annotated[AppState, Depends(dependency=AppState)],
@@ -77,7 +99,7 @@ def create_app() -> FastAPI:
 
     @app.get(path="/get_huggingface_video")
     def endpoint_get_huggingface_video(
-        env_id: SingleAgentId,
+        env_id: GymId,
         hf_client: Annotated[HuggingFaceClient, Depends(dependency=HuggingFaceClient)],
     ) -> StreamingResponse:
         hf_client.mp4_paths[env_id]
@@ -85,7 +107,7 @@ def create_app() -> FastAPI:
 
     @app.get(path="/get_env_video")
     def endpoint_get_env_video(
-        env_id: SingleAgentId,
+        env_id: GymId,
     ) -> StreamingResponse:
         env = make(id=env_id, render_mode="rgb_array")
         env = RecordVideo(
@@ -93,29 +115,30 @@ def create_app() -> FastAPI:
             video_folder="tmp",
         )
         env.reset(seed=123)
+
+        if env.video_recorder is None:
+            msg = "env.video_recorder is None"
+            raise ValueError(msg)
+
         agent = RandomAgent[Any, Any]()
         terminated, truncated = False, False
         while not (terminated or truncated):
-            action = agent.get_action(env=env)
+            action = agent.get_action(env=env)  # type: ignore[arg-type]
             _, _, terminated, truncated, _ = env.step(action=action)
             env.render()
         env.video_recorder.close()
         return stream_mp4(mp4_path=Path(env.video_recorder.path))
 
     @app.exception_handler(exc_class_or_status_code=RequestValidationError)
-    async def validation_exception_handler(
-        request: Request, exc: RequestValidationError
-    ) -> JSONResponse:
+    async def validation_exception_handler(request: Request, exc: RequestValidationError) -> JSONResponse:
         logger.debug(f"url: {request.url}")
         if hasattr(request, "_body"):
-            logger.debug(f"body: {request._body}")
+            logger.debug(f"body: {request._body.decode()}")  # noqa: SLF001
         logger.debug(f"header: {request.headers}")
         logger.error(f"{request}: {exc}")
-        exc_str = f"{exc}".replace("\n", " ").replace("   ", " ")
+        exc_str = str(exc).replace("\n", " ").replace("   ", " ")
         content = {"status_code": 10422, "message": exc_str, "data": None}
-        return JSONResponse(
-            content=content, status_code=status.HTTP_422_UNPROCESSABLE_ENTITY
-        )
+        return JSONResponse(content=content, status_code=status.HTTP_422_UNPROCESSABLE_ENTITY)
 
     app.add_middleware(
         middleware_class=CORSMiddleware,
@@ -125,17 +148,3 @@ def create_app() -> FastAPI:
         allow_headers=["*"],
     )
     return app
-
-if __name__ == "__main__":
-    import uvicorn
-    import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--host", type=str, default="0.0.0.0")
-    parser.add_argument("--port", type=int, default=7860)
-    parser.add_argument("--log", type=str, default="info")
-    parser.add_argument('--reload', action='store_true', help='Reload flag')
-    args = parser.parse_args()
-
-    config = uvicorn.Config(app=create_app(), host=args.host, port=args.port, log_level=args.log, reload=args.reload)
-    server = uvicorn.Server(config=config)
-    server.run()

src/app_state.py

@@ -1,45 +1,40 @@
 from typing import Any, Self
 
+from loguru import logger
+from src.typing import AgentType, CpuConfig, RolloutPolicy
+
 from litrl import make_multiagent
 from litrl.algo.mcts.agent import MCTSAgent
-from litrl.algo.mcts.mcts_config import MCTSConfig
+from litrl.algo.mcts.mcts_config import MCTSConfigBuilder
 from litrl.algo.mcts.rollout import VanillaRollout
-from litrl.algo.sac.agent import OnnxSacDeterministicAgent
 from litrl.common.agent import Agent, RandomMultiAgent
-from litrl.env.connect_four import Board
-from litrl.env.set_state import set_state
-from loguru import logger
-from src.typing import AgentType, CpuConfig, RolloutPolicy
 from litrl.env.connect_four import ConnectFour
+from litrl.model.sac.multi_agent import OnnxSacDeterministicMultiAgent
+
 
 class AppState:
     _instance: Self | None = None
     env: ConnectFour
     cpu_config: CpuConfig
-    agent: Agent[Any, Any]
+    agent: Agent[Any, int]
 
     def setup(self) -> None:
         logger.debug("AppState setup called")
-        self.env = make_multiagent(id="ConnectFour-v3", render_mode="rgb_array")
+        self.env = make_multiagent(id="connect_four", render_mode="rgb_array")
         self.env.reset(seed=123)
 
         self.cpu_config: CpuConfig = CpuConfig(agent_type=AgentType.RANDOM)
         self.set_agent()  # TODO in properties setter.
-        self.agent: Agent[Any, Any]
+        self.agent: Agent[Any, int]
 
-    def __new__(cls):
+    def __new__(cls) -> "AppState":
         if cls._instance is None:
             cls._instance = super().__new__(cls)
             cls._instance.setup()
         return cls._instance
-    
-    def set_board(self, board: Board) -> None:
-        set_state(env=self.env, board=board)
 
     def set_config(self, cpu_config: CpuConfig) -> None:
-        if (
-            cpu_config != self.cpu_config
-        ):
+        if cpu_config != self.cpu_config:
             self.cpu_config = cpu_config
             self.set_agent()
 
@@ -48,31 +43,32 @@ class AppState:
             case None:
                 return RandomMultiAgent()
             case RolloutPolicy.SAC:
-                return OnnxSacDeterministicAgent()
+                return OnnxSacDeterministicMultiAgent()
             case RolloutPolicy.RANDOM:
                 return RandomMultiAgent()
             case _:
-                raise NotImplementedError(
-                    f"cpu_config.rollout_policy: {self.cpu_config.rollout_policy}"
-                )
+                msg = f"cpu_config.rollout_policy: {self.cpu_config.rollout_policy}"
+                raise NotImplementedError(msg)
 
     def set_agent(self) -> None:
         match self.cpu_config.agent_type:
             case AgentType.MCTS:
                 rollout_agent = self.create_rollout()
-                mcts_config = MCTSConfig(
-                    simulations=self.cpu_config.simulations or 50,
-                    rollout_strategy=VanillaRollout(rollout_agent=rollout_agent),
-                )
+                # fmt: off
+                mcts_config = (
+                    MCTSConfigBuilder()
+                    .set_simulations(self.cpu_config.simulations or 50)
+                    .set_rollout_strategy(VanillaRollout(rollout_agent=rollout_agent))
+                ).build()
+                # fmt: on
                 self.agent = MCTSAgent(cfg=mcts_config)
             case AgentType.RANDOM:
                 self.agent = RandomMultiAgent()
             case AgentType.SAC:
-                self.agent = OnnxSacDeterministicAgent()
+                self.agent = OnnxSacDeterministicMultiAgent()  # type: ignore[assignment]  # TODO
             case _:
-                raise NotImplementedError(
-                    f"cpu_config.name: {self.cpu_config.agent_type}"
-                )
+                msg = f"cpu_config.name: {self.cpu_config.agent_type}"
+                raise NotImplementedError(msg)
 
     def get_action(self) -> int:
         return self.agent.get_action(env=self.env)

src/constants.py

@@ -1,6 +1,4 @@
-SPACE_REPO = "c-gohlke/litrl"
-SPACE_REPO_TYPE = "space"
 MODEL_REPO = "c-gohlke/litrl"
 MODEL_REPO_TYPE = "model"
 
-ENV_RESULTS_FILE_DEPTH = 3
+ENV_RESULTS_FILE_DEPTH = 3

src/huggingface/get_environments.py

@@ -1,5 +1,6 @@
-from huggingface_hub import HfApi  # type: ignore[import]
-from src.constants import MODEL_REPO, MODEL_REPO_TYPE, ENV_RESULTS_FILE_DEPTH
+from huggingface_hub import HfApi
+from src.constants import ENV_RESULTS_FILE_DEPTH, MODEL_REPO, MODEL_REPO_TYPE
+
 
 def get_environments(hf_api: HfApi) -> list[str]:
     environments: list[str] = []
@@ -9,4 +10,4 @@ def get_environments(hf_api: HfApi) -> list[str]:
         # e.g. ['models', 'CartPole-v1', 'results.yaml']
         if len(vals) == ENV_RESULTS_FILE_DEPTH and vals[2] == "results.yaml" and vals[0] == "models":
             environments.append(vals[1])
-    return environments
+    return environments

src/huggingface/get_files.py

@@ -1,8 +1,10 @@
 from pathlib import Path
-from huggingface_hub import hf_hub_download  # type: ignore[import]
-import yaml  # type: ignore[import]
+
+import yaml
+from huggingface_hub import hf_hub_download
 from src.constants import MODEL_REPO, MODEL_REPO_TYPE
 
+
 def get_mp4_paths(environments: list[str]) -> dict[str, Path]:
     mp4_paths: dict[str, Path] = {}
     for env in environments:
@@ -12,15 +14,23 @@ def get_mp4_paths(environments: list[str]) -> dict[str, Path]:
 
 def get_best(env: str, filename: str = "demo.mp4") -> Path:
     hf_env_results_path = f"models/{env}/results.yaml"
-    local_env_results_path = hf_hub_download(repo_id=MODEL_REPO, repo_type=MODEL_REPO_TYPE, filename=hf_env_results_path)
-    with open(local_env_results_path, "r") as f:
+    local_env_results_path = hf_hub_download(
+        repo_id=MODEL_REPO,
+        repo_type=MODEL_REPO_TYPE,
+        filename=hf_env_results_path,
+    )
+    with Path(local_env_results_path).open() as f:
         env_results = yaml.load(f, Loader=yaml.FullLoader)
     best_model_type = max(env_results, key=lambda model: env_results[model])
     model_results_path = f"models/{env}/{best_model_type}/results.yaml"
-    local_model_results_path = hf_hub_download(repo_id=MODEL_REPO, repo_type=MODEL_REPO_TYPE, filename=model_results_path)
-    with open(local_model_results_path, "r") as f:
+    local_model_results_path = hf_hub_download(
+        repo_id=MODEL_REPO,
+        repo_type=MODEL_REPO_TYPE,
+        filename=model_results_path,
+    )
+    with Path(local_model_results_path).open() as f:
         model_results = yaml.load(f, Loader=yaml.FullLoader)
 
     best_model = max(model_results, key=lambda model: model_results[model])
     hf_gif_path = f"models/{env}/{best_model_type}/{best_model}/{filename}"
-    return Path(hf_hub_download(repo_id=MODEL_REPO, repo_type=MODEL_REPO_TYPE, filename=hf_gif_path))
+    return Path(hf_hub_download(repo_id=MODEL_REPO, repo_type=MODEL_REPO_TYPE, filename=hf_gif_path))

src/huggingface/huggingface_client.py

@@ -1,13 +1,14 @@
-from huggingface_hub import HfApi, login  # type: ignore[import]
 import os
-from loguru import logger
-from .get_files import get_mp4_paths
+
+from huggingface_hub import HfApi, login
+
 from .get_environments import get_environments
+from .get_files import get_mp4_paths
 
 
 class HuggingFaceClient:
     def __init__(self) -> None:
-        login(  # type: ignore[no-untyped-call]
+        login(
             token=os.environ.get("HUGGINGFACE_TOKEN"),
             add_to_git_credential=True,
             new_session=False,
@@ -15,11 +16,3 @@ class HuggingFaceClient:
         self.hf_api = HfApi()
         self.environments = get_environments(self.hf_api)
         self.mp4_paths = get_mp4_paths(environments=self.environments)
-        
-
-    def api_predict(self, env_id: str)-> bytes|None:
-        if env_id not in self.mp4_paths:
-            logger.error(f"Environment {env_id} not found in {self.mp4_paths}")
-            return None
-        with open(self.mp4_paths[env_id], "rb") as f:
-            return f.read()

src/typing.py

@@ -17,4 +17,4 @@ class RolloutPolicy(enum.Enum):
 class CpuConfig(BaseModel):
     agent_type: AgentType
     simulations: int | None = None
-    rollout_policy: RolloutPolicy | None = None
+    rollout_policy: RolloutPolicy | None = None