| from collections import defaultdict, namedtuple |
| import math |
| from pathlib import Path |
| from typing import Any, Dict, List, Tuple |
|
|
| import pygame |
| import torch |
| from torch import Tensor |
|
|
| from agent import Agent |
| from data import Dataset, Episode |
| from game.keymap import ActionNames, Keymap |
| from envs import WorldModelEnv |
|
|
|
|
| NamedEnv = namedtuple("NamedEnv", "name env") |
| OneStepData = namedtuple("OneStepData", "obs act rew end trunc") |
|
|
|
|
| class PlayEnv: |
| def __init__( |
| self, |
| agent: Agent, |
| envs: List[NamedEnv], |
| action_names: ActionNames, |
| keymap: Keymap, |
| recording_mode: bool, |
| store_denoising_trajectory: bool, |
| store_original_obs: bool, |
| ) -> None: |
| self.agent = agent |
| self.envs = envs |
| self.action_names = action_names |
| self.keymap = keymap |
| self.recording_mode = recording_mode |
| self.store_denoising_trajectory = store_denoising_trajectory |
| self.store_original_obs = store_original_obs |
| self.is_human_player = False |
| self.env_id = 0 |
| self.env_name, self.env = self.envs[0] |
| self.obs, self.t, self.return_, self.hx_cx, self.ckpt_id, self.buffer, self.rec_dataset = (None,) * 7 |
|
|
| def print_controls(self) -> None: |
| print("\nControls (play mode):\n") |
| print("m : controller (policy/human)") |
| print("↑ : imagination horizon (+1)") |
| print("↓ : imagination horizon (-1)") |
| print(f"→ : next environment ({' → '.join([env_name for (env_name, _) in self.envs])})") |
| print(f"← : prev environment ({' ← '.join([env_name for (env_name, _) in self.envs])})") |
| print("\nEnvironment actions:\n") |
| for keys, idx in self.keymap.items(): |
| key_names = [pygame.key.name(key) for key in keys] |
| key_names = ["⎵" if key_name == "space" else key_name for key_name in key_names] |
| print(f"{' + '.join(key_names)} : {self.action_names[idx]}") |
|
|
| def next_mode(self) -> bool: |
| self.switch_controller() |
| return True |
|
|
| def next_axis_1(self) -> bool: |
| self.update_wm_horizon(+1) |
| return True |
|
|
| def prev_axis_1(self) -> bool: |
| self.update_wm_horizon(-1) |
| return True |
|
|
| def next_axis_2(self) -> bool: |
| self.switch_env(self.env_id + 1) |
| return True |
|
|
| def prev_axis_2(self) -> bool: |
| self.switch_env(self.env_id - 1) |
| return True |
|
|
| def is_wm_env(self) -> bool: |
| return isinstance(self.env, WorldModelEnv) |
|
|
| def print_env(self) -> None: |
| print(f"> Environment: {self.env_name}") |
|
|
| def print_control(self) -> None: |
| print(f"> Control: {'human' if self.is_human_player else 'policy'}") |
|
|
| def switch_env(self, env_id: int) -> None: |
| self.env_id = env_id % len(self.envs) |
| self.env_name, self.env = self.envs[self.env_id] |
| self.print_env() |
|
|
| def switch_controller(self) -> None: |
| self.is_human_player = not self.is_human_player |
| self.print_control() |
|
|
| def update_wm_horizon(self, incr: int) -> None: |
| if self.is_wm_env(): |
| self.env.horizon = max(1, self.env.horizon + incr) |
|
|
| def reset_recording(self) -> None: |
| self.buffer = defaultdict(list) |
| self.buffer["info"] = defaultdict(list) |
| dir = Path("dataset") / f"rec_{self.env_name}_{'H' if self.is_human_player else 'π'}" |
| self.rec_dataset = Dataset(dir) |
| self.rec_dataset.load_from_default_path() |
|
|
| def reset(self) -> Tuple[Tensor, None]: |
| self.obs, _ = self.env.reset() |
| self.t, self.return_, self.hx_cx = 0, 0, None |
| if self.recording_mode: |
| self.reset_recording() |
| return self.obs, None |
|
|
| @torch.no_grad() |
| def step(self, act: int) -> Tuple[Tensor, Tensor, Tensor, Tensor, Dict[str, Any]]: |
| if self.is_human_player: |
| act = torch.tensor([act], device=self.agent.device) |
| else: |
| logits_act, value, self.hx_cx = self.agent.actor_critic.predict_act_value(self.obs, self.hx_cx) |
| dst = torch.distributions.categorical.Categorical(logits=logits_act) |
| act = dst.sample() |
| entropy = dst.entropy() / math.log(2) |
| entropy = None if self.is_human_player else f"{entropy.item():.2f}" |
| value = None if self.is_human_player else f"{value.item():.2f}" |
| next_obs, rew, end, trunc, env_info = self.env.step(act) |
| data = OneStepData(self.obs, act, rew, end, trunc) |
| self.return_ += rew.item() |
| control = "human" if self.is_human_player else "policy" |
| header = [ |
| [ |
| f"Env : {self.env_name}", |
| f"Control : {control}", |
| f"Timestep: {self.t + 1}", |
| f"Horizon : {self.env.horizon}" if self.is_wm_env() else "", |
| ], |
| [ |
| f"Trunc : {bool(trunc)}", |
| f"Done : {bool(end)}", |
| f"Reward: {rew.item():.2f}", |
| f"Return: {self.return_:.2f}", |
| ], |
| [ |
| f"Action : {self.action_names[act[0]]}", |
| f"Entropy: {entropy}", |
| f"Value : {value}", |
| ], |
| ] |
| info = {"header": header} |
|
|
| if end or trunc: |
| d = "Dead" if end else ("Horizon" if self.is_wm_env() else "Timed out") |
| print(f"[{d}] return = {self.return_} - length = {self.t + 1}.") |
|
|
| if self.recording_mode: |
| for k, v in data._asdict().items(): |
| self.buffer[k].append(v) |
| if self.store_denoising_trajectory and "denoising_trajectory" in env_info: |
| self.buffer["info"]["denoising_trajectory"].append(env_info["denoising_trajectory"]) |
| if self.store_original_obs and "original_obs" in env_info: |
| original_obs = (torch.tensor(env_info["original_obs"][0]).permute(2, 0, 1).unsqueeze(0).contiguous()) |
| self.buffer["info"]["original_obs"].append(original_obs) |
| if end or trunc: |
| ep_dict = {k: torch.cat(v, dim=0) for k, v in self.buffer.items() if k != "info"} |
| ep_info = {k: torch.cat(v, dim=0) for k, v in self.buffer["info"].items()} |
| ep = Episode(**ep_dict, info=ep_info).to("cpu") |
| self.rec_dataset.add_episode(ep) |
| self.rec_dataset.save_to_default_path() |
|
|
| self.obs = next_obs |
| self.t += 1 |
|
|
| return next_obs, rew, end, trunc, info |
|
|
