Tutorial for enn-ppo (#230)

Author: cswinter

enn_ppo/enn_ppo/__init__.py

@@ -0,0 +1,6 @@
+from .load_checkpoint import load_agent, load_checkpoint
+
+__all__ = [
+    "load_checkpoint",
+    "load_agent",
+]

enn_ppo/enn_ppo/agent.py

@@ -5,8 +5,13 @@
 import torch
 from ragged_buffer import RaggedBufferBool, RaggedBufferF32, RaggedBufferI64
 
-from entity_gym.environment import VecActionMask
+import entity_gym.agent
+from entity_gym.environment import Action, Observation, VecActionMask
+from entity_gym.environment.env_list import action_index_to_actions
+from entity_gym.environment.environment import ActionType
+from entity_gym.environment.vec_env import batch_obs
 from entity_gym.simple_trace import Tracer
+from rogue_net.rogue_net import RogueNet
 
 
 class PPOAgent(Protocol):
@@ -41,3 +46,26 @@ def get_auxiliary_head(
         tracer: Tracer,
     ) -> torch.Tensor:
         ...
+
+
+class RogueNetAgent(entity_gym.agent.Agent):
+    def __init__(self, agent: RogueNet):
+        self.agent = agent
+
+    def act(self, obs: Observation) -> Tuple[Dict[ActionType, Action], float]:
+        vec_obs = batch_obs([obs], self.agent.obs_space, self.agent.action_space)
+        with torch.no_grad():
+            act_indices, _, _, _, aux, logits = self.agent.get_action_and_auxiliary(
+                vec_obs.features,
+                vec_obs.visible,
+                vec_obs.action_masks,
+                tracer=Tracer(False),
+            )
+        actions = action_index_to_actions(
+            self.agent.obs_space,
+            self.agent.action_space,
+            act_indices,
+            obs,
+            probs={k: l.exp().cpu().numpy() for k, l in logits.items()},
+        )
+        return actions, float(aux["value"].item())

enn_ppo/enn_ppo/load_checkpoint.py

@@ -0,0 +1,12 @@
+from hyperstate import StateManager
+
+from .config import TrainConfig
+from .train import State, initialize
+
+
+def load_checkpoint(path: str) -> StateManager[TrainConfig, State]:
+    return StateManager(TrainConfig, State, initialize, init_path=path)
+
+
+def load_agent(path: str) -> State:
+    return StateManager(TrainConfig, State, initialize, init_path=path).state

enn_ppo/enn_ppo/train.py

@@ -8,6 +8,7 @@
 from pathlib import Path
 from typing import Any, Callable, Dict, Mapping, Optional, Type, Union
 
+import click
 import hyperstate
 import numpy as np
 import torch
@@ -42,12 +43,10 @@ def serialize(self) -> Any:
     def deserialize(
         clz, state_dict: Any, config: TrainConfig, state: "State", ctx: Dict[str, Any]
     ) -> "SerializableRogueNet":
-        obs_space: ObsSpace = ctx["obs_space"]
-        action_space: Dict[ActionType, ActionSpace] = ctx["action_space"]
         net = SerializableRogueNet(
             config.net,
-            obs_space,
-            action_space,
+            state.obs_space,
+            state.action_space,
             regression_heads={"value": 1},
         )
         net.load_state_dict(state_dict)
@@ -114,6 +113,8 @@ class State(hyperstate.Lazy):
     value_function: Optional[SerializableRogueNet]
     optimizer: SerializableAdamW
     vf_optimizer: Optional[SerializableAdamW]
+    obs_space: ObsSpace
+    action_space: Dict[str, ActionSpace]
 
 
 def train(
@@ -358,26 +359,6 @@ def _run_eval() -> None:
                 writer.add_scalar(f"{name}.max", value.max, global_step)
                 writer.add_scalar(f"{name}.min", value.min, global_step)
                 writer.add_scalar(f"{name}.count", value.count, global_step)
-            # Double log these to remain compatible with old naming scheme
-            # TODO: remove before release
-            writer.add_scalar(
-                "charts/episodic_return",
-                metrics["episodic_reward"].mean,
-                global_step,
-            )
-            writer.add_scalar(
-                "charts/episodic_length",
-                metrics["episode_length"].mean,
-                global_step,
-            )
-            writer.add_scalar(
-                "charts/episodes", metrics["episodic_reward"].count, global_step
-            )
-            writer.add_scalar("meanrew", metrics["reward"].mean, global_step)
-
-        print(
-            f"global_step={global_step} {'  '.join(f'{name}={value.mean}' for name, value in metrics.items())}"
-        )
 
         values = rollout.values
         actions = rollout.actions
@@ -571,9 +552,41 @@ def _run_eval() -> None:
                                 np.sum(_actions == i).item() / len(_actions),
                                 global_step,
                             )
-            print(
-                "SPS:", int((global_step - initial_step) / (time.time() - start_time))
+
+            fps = (global_step - initial_step) / (time.time() - start_time)
+            digits = int(np.ceil(np.log10(cfg.total_timesteps)))
+            episodic_reward = metrics["episodic_reward"].mean
+            episode_length = metrics["episode_length"].mean
+            episode_count = metrics["episode_length"].count
+            mean_reward = metrics["reward"].mean
+
+            def green(s: str) -> str:
+                return click.style(s, fg="cyan")
+
+            def estyle(f: float) -> str:
+                return click.style(f"{f:.2e}", fg="cyan")
+
+            def fstyle(f: float) -> str:
+                return click.style(f"{f:5.2f}", fg="cyan")
+
+            def tstyle(s: str) -> str:
+                return s
+
+            def symstyle(s: str) -> str:
+                return click.style(s, fg="white", bold=True)
+
+            # fmt: off
+            click.echo(
+                green(f"{global_step:>{digits}}") + symstyle("/") + green(f"{cfg.total_timesteps} ")
+                + f"{symstyle('|')} {tstyle('meanrew')} {estyle(mean_reward)} "
+                + f"{symstyle('|')} {tstyle('explained_var')} {fstyle(explained_var.item())} "
+                + f"{symstyle('|')} {tstyle('entropy')} {fstyle(entropy_loss.item())} "
+                + f"{symstyle('|')} {tstyle('episodic_reward')} {estyle(episodic_reward)} "
+                + f"{symstyle('|')} {tstyle('episode_length')} {estyle(episode_length)} "
+                + f"{symstyle('|')} {tstyle('episodes')} {green(str(episode_count))} "
+                + f"{symstyle('|')} {tstyle('fps')} {green(str(int(fps)))}"
             )
+            # fmt: on
             writer.add_scalar(
                 "charts/SPS",
                 int((global_step - initial_step) / (time.time() - start_time)),
@@ -679,6 +692,8 @@ def initialize(cfg: TrainConfig, ctx: Dict[str, Any]) -> State:
         value_function=value_function,
         optimizer=optimizer,
         vf_optimizer=vf_optimizer,
+        obs_space=ctx["obs_space"],
+        action_space=ctx["action_space"],
     )