Upload folder using huggingface_hub

src/app.py

@@ -1,6 +1,7 @@
 from typing import Annotated, Any, 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
@@ -13,7 +14,9 @@ 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.huggingface.huggingface_client import HuggingFaceClient
+from litrl.env.connect_four import ConnectFour
 
 def stream_mp4(mp4_path: Path) -> StreamingResponse:
     def iter_file()-> Generator[bytes, Any, None]:  
@@ -22,6 +25,19 @@ def stream_mp4(mp4_path: Path) -> StreamingResponse:
 
     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:
+    env.step(action)
+    return observe(env)
+
+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}
 
 def create_app() -> FastAPI:
     app = FastAPI()
@@ -30,44 +46,45 @@ def create_app() -> FastAPI:
     async def to_docs():
         return RedirectResponse("/docs")
 
-    @app.post("/", response_model=int)
-    def bot_action(
-        board: Board,
-        cpuConfig: CpuConfig,
+    @app.post(path="/connect_four/play", response_model=GridResponseType)
+    def endpoint_play(
+        action: int,
+        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)],
-    ) -> int:
-        app_state.set_config(cpu_config=cpuConfig)
-        app_state.set_board(board=board)
-        return app_state.get_action()
+    ) -> GridResponseType:
+        return observe(app_state.env)
+    
+    @app.post(path="/connect_four/bot_play", response_model=GridResponseType)
+    def endpoint_bot_play(
+        cpu_config: CpuConfig,
+        app_state: Annotated[AppState, Depends(dependency=AppState)],
+    ) -> GridResponseType:
+        app_state.set_config(cpu_config)
+        action = app_state.get_action()
+        return step(app_state.env, action)
 
-    @app.post(path=f"/game", response_model=str)
-    def bot_action(
-        env_id: SingleAgentId,
-    ) -> str:
-        env = RecordVideo(
-            env=make(id=env_id, render_mode="rgb_array"),
-            video_folder="tmp",
-        )
-        env.reset(seed=123)
-        agent = RandomAgent[Any, Any]()
-        terminated, truncated = False, False
-        while not (terminated or truncated):
-            action = agent.get_action(env=env)
-            _, _, terminated, truncated, _ = env.step(action=action)
-            env.render()
-        env.video_recorder.close()
-        return stream_mp4(mp4_path=Path(env.video_recorder.path))
+    @app.get(path="/connect_four/reset", response_model=GridResponseType)
+    def endpoint_reset(
+        app_state: Annotated[AppState, Depends(dependency=AppState)],
+    ) -> GridResponseType:
+        app_state.env.reset()
+        return observe(app_state.env)
 
-    @app.get(path=f"/hfmp4")
-    def fh_stream(
+    @app.get(path="/get_huggingface_video")
+    def endpoint_get_huggingface_video(
         env_id: SingleAgentId,
         hf_client: Annotated[HuggingFaceClient, Depends(dependency=HuggingFaceClient)],
     ) -> StreamingResponse:
         hf_client.mp4_paths[env_id]
         return stream_mp4(mp4_path=hf_client.mp4_paths[env_id])
 
-    @app.get(path=f"/mp4")
-    def bot_action(
+    @app.get(path="/get_env_video")
+    def endpoint_get_env_video(
         env_id: SingleAgentId,
     ) -> StreamingResponse:
         env = make(id=env_id, render_mode="rgb_array")
@@ -115,7 +132,10 @@ if __name__ == "__main__":
     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="info")
+
+    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,4 +1,4 @@
-from typing import Any
+from typing import Any, Self
 
 from litrl import make_multiagent
 from litrl.algo.mcts.agent import MCTSAgent
@@ -8,32 +8,42 @@ 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
 
 class AppState:
-    def __init__(self) -> None:
-        self.env = make_multiagent(id="ConnectFour-v3", render_mode="human")
+    _instance: Self | None = None
+    env: ConnectFour
+    cpu_config: CpuConfig
+    agent: Agent[Any, Any]
+
+    def setup(self) -> None:
+        logger.debug("AppState setup called")
+        self.env = make_multiagent(id="ConnectFour-v3", render_mode="rgb_array")
         self.env.reset(seed=123)
-        self.agent: Agent[Any, Any] | None = None
-        self.cpu_config: CpuConfig | None = None
 
+        self.cpu_config: CpuConfig = CpuConfig(agent_type=AgentType.RANDOM)
+        self.set_agent()  # TODO in properties setter.
+        self.agent: Agent[Any, Any]
+
+    def __new__(cls):
+        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 (
-            self.agent is None
-            or self.cpu_config is None
-            or cpu_config != self.cpu_config
+            cpu_config != self.cpu_config
         ):
             self.cpu_config = cpu_config
             self.set_agent()
 
     def create_rollout(self) -> Agent[Any, Any]:
-        if self.cpu_config is None:
-            raise ValueError("self.cpu_config is None")
         match self.cpu_config.rollout_policy:
             case None:
                 return RandomMultiAgent()
@@ -47,9 +57,6 @@ class AppState:
                 )
 
     def set_agent(self) -> None:
-        if self.cpu_config is None:
-            raise ValueError("self.cpu_config is None")
-
         match self.cpu_config.agent_type:
             case AgentType.MCTS:
                 rollout_agent = self.create_rollout()
@@ -68,6 +75,4 @@ class AppState:
                 )
 
     def get_action(self) -> int:
-        if self.agent is None:
-            raise ValueError("self.agent is None")
         return self.agent.get_action(env=self.env)