Add mechanism to save frames when visualizing. Add option to visualize agent playing in the world model.

src/envs/world_model_env.py

@@ -9,24 +9,15 @@ import torch
 from torch.distributions.categorical import Categorical
 import torchvision
 
-from utils import extract_state_dict
-
 
 class WorldModelEnv:
 
-    def __init__(self, tokenizer: torch.nn.Module, world_model: torch.nn.Module, device: Union[str, torch.device], pretrained_agent_path: Optional[str] = None, env: Optional[gym.Env] = None) -> None:
+    def __init__(self, tokenizer: torch.nn.Module, world_model: torch.nn.Module, device: Union[str, torch.device], env: Optional[gym.Env] = None) -> None:
 
         self.device = torch.device(device)
         self.world_model = world_model.to(self.device).eval()
         self.tokenizer = tokenizer.to(self.device).eval()
 
-        if pretrained_agent_path is not None:
-            agent_state_dict = torch.load(pretrained_agent_path)
-            self.world_model.load_state_dict(extract_state_dict(agent_state_dict, 'world_model'))
-            incompatible_keys = self.tokenizer.load_state_dict(extract_state_dict(agent_state_dict, 'tokenizer'), strict=False)
-            assert not incompatible_keys.missing_keys
-            assert (not incompatible_keys.unexpected_keys) or all([k.startswith('lpips.') for k in incompatible_keys.unexpected_keys])
-
         self.keys_values_wm, self.obs_tokens, self._num_observations_tokens = None, None, None
 
         self.env = env

src/envs/wrappers.py

@@ -30,7 +30,7 @@ class ResizeObsWrapper(gym.ObservationWrapper):
         gym.ObservationWrapper.__init__(self, env)
         self.size = tuple(size)
         self.observation_space = gym.spaces.Box(low=0, high=255, shape=(size[0], size[1], 3), dtype=np.uint8)
-        self.original_obs = None
+        self.unwrapped.original_obs = None
 
     def resize(self, obs: np.ndarray):
         img = Image.fromarray(obs)
@@ -38,7 +38,7 @@ class ResizeObsWrapper(gym.ObservationWrapper):
         return np.array(img)
 
     def observation(self, observation: np.ndarray) -> np.ndarray:
-        self.original_obs = observation
+        self.unwrapped.original_obs = observation
         return self.resize(observation)
 
 

src/game/agent_env.py

@@ -2,23 +2,26 @@ from einops import rearrange
 import numpy as np
 from PIL import Image
 import torch
+from torchvision.transforms.functional import InterpolationMode, resize
 
 from agent import Agent
-from envs import SingleProcessEnv
+from envs import SingleProcessEnv, WorldModelEnv
 from game.keymap import get_keymap_and_action_names
 
 
 class AgentEnv:
-    def __init__(self, agent: Agent, env: SingleProcessEnv, keymap_name: str) -> None:
-        assert isinstance(env, SingleProcessEnv)
+    def __init__(self, agent: Agent, env: SingleProcessEnv, keymap_name: str, do_reconstruction: bool) -> None:
+        assert isinstance(env, SingleProcessEnv) or isinstance(env, WorldModelEnv)
         self.agent = agent
         self.env = env
         _, self.action_names = get_keymap_and_action_names(keymap_name)
+        self.do_reconstruction = do_reconstruction
         self.obs = None
         self._t = None
         self._return = None
 
     def _to_tensor(self, obs: np.ndarray):
+        assert isinstance(obs, np.ndarray) and obs.dtype == np.uint8
         return rearrange(torch.FloatTensor(obs).div(255), 'n h w c -> n c h w').to(self.agent.device)
 
     def _to_array(self, obs: torch.FloatTensor):
@@ -27,7 +30,7 @@ class AgentEnv:
 
     def reset(self):
         obs = self.env.reset()
-        self.obs = self._to_tensor(obs)
+        self.obs = self._to_tensor(obs) if isinstance(self.env, SingleProcessEnv) else obs
         self.agent.actor_critic.reset(1)
         self._t = 0
         self._return = 0
@@ -37,7 +40,7 @@ class AgentEnv:
         with torch.no_grad():
             act = self.agent.act(self.obs, should_sample=True).cpu().numpy()
         obs, reward, done, _ = self.env.step(act)
-        self.obs = self._to_tensor(obs)
+        self.obs = self._to_tensor(obs) if isinstance(self.env, SingleProcessEnv) else obs
         self._t += 1
         self._return += reward[0]
         info = {
@@ -49,6 +52,13 @@ class AgentEnv:
 
     def render(self) -> Image.Image:
         assert self.obs.size() == (1, 3, 64, 64)
-        rec = torch.clamp(self.agent.tokenizer.encode_decode(self.obs, should_preprocess=True, should_postprocess=True), 0, 1)
-        arr = self._to_array(torch.cat((self.obs, rec), dim=-1))
+        original_obs = self.env.env.unwrapped.original_obs if isinstance(self.env, SingleProcessEnv) else self._to_array(self.obs)
+        if self.do_reconstruction:
+            rec = torch.clamp(self.agent.tokenizer.encode_decode(self.obs, should_preprocess=True, should_postprocess=True), 0, 1)
+            rec = self._to_array(resize(rec, original_obs.shape[:2], interpolation=InterpolationMode.NEAREST))
+            resized_obs = self._to_array(resize(self.obs, original_obs.shape[:2], interpolation=InterpolationMode.NEAREST))
+            arr = np.concatenate((original_obs, resized_obs, rec), axis=1)
+        else:
+            arr = original_obs
         return Image.fromarray(arr)
+

src/game/game.py

@@ -1,3 +1,5 @@
+from datetime import datetime
+from pathlib import Path
 from typing import Tuple, Union
 
 import gym
@@ -7,17 +9,20 @@ from PIL import Image
 
 from envs import WorldModelEnv
 from game.keymap import get_keymap_and_action_names
+from utils import make_video
 
 
 class Game:
-    def __init__(self, env: Union[gym.Env, WorldModelEnv], keymap_name: str, size: Tuple[int, int], fps: int, verbose: bool) -> None:
+    def __init__(self, env: Union[gym.Env, WorldModelEnv], keymap_name: str, size: Tuple[int, int], fps: int, verbose: bool, record_mode: bool) -> None:
         self.env = env
         self.height, self.width = size
         self.fps = fps
         self.verbose = verbose
-
+        self.record_mode = record_mode
         self.keymap, self.action_names = get_keymap_and_action_names(keymap_name)
 
+        self.record_dir = Path('media') / 'recordings'
+
         print('Actions:')
         for key, idx in self.keymap.items():
             print(f'{pygame.key.name(key)}: {self.action_names[idx]}')
@@ -52,16 +57,23 @@ class Game:
 
         if isinstance(self.env, gym.Env):
             _, info = self.env.reset(return_info=True)
-            draw_game(info['rgb'])
+            img = info['rgb']
         else:
             self.env.reset()
-            draw_game(self.env.render())
+            img = self.env.render()
+        
+        draw_game(img)
 
         clear_header()
         pygame.display.flip()
 
+        episode_buffer = []
+        segment_buffer = []
+        recording = False
+
         do_reset, do_wait = False, False
         should_stop = False
+        
         while not should_stop:
 
             action = 0  # noop
@@ -75,6 +87,16 @@ class Game:
                     do_reset = True
                 if event.type == pygame.KEYDOWN and event.key == pygame.K_PERIOD:
                     do_wait = not do_wait
+                if event.type == pygame.KEYDOWN and event.key == pygame.K_COMMA:
+                    if not recording: 
+                        recording = True
+                        print('Started recording.')
+                    else:
+                        print('Stopped recording.')    
+                        self.save_recording(np.stack(segment_buffer))
+                        recording = False
+                        segment_buffer = []
+
             if action == 0:
                 pressed = pygame.key.get_pressed()
                 for key, action in self.keymap.items():
@@ -91,6 +113,12 @@ class Game:
             img = info['rgb'] if isinstance(self.env, gym.Env) else self.env.render()
             draw_game(img)
 
+            if recording:
+                segment_buffer.append(np.array(img))
+
+            if self.record_mode:
+                episode_buffer.append(np.array(img))
+
             if self.verbose:
                 clear_header()
                 draw_text(f'Action: {self.action_names[action]}', idx_line=0)
@@ -108,4 +136,16 @@ class Game:
                 self.env.reset()
                 do_reset = False
 
+                if self.record_mode:
+                    if input('Save episode? [Y/n] ').lower() != 'n':
+                        self.save_recording(np.stack(episode_buffer))
+                    episode_buffer = []
+
         pygame.quit()
+
+    def save_recording(self, frames):
+        self.record_dir.mkdir(exist_ok=True, parents=True)
+        timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+        np.save(self.record_dir / timestamp, frames)
+        make_video(self.record_dir / f'{timestamp}.mp4', fps=15, frames=frames)
+        print(f'Saved recording {timestamp}.')

src/play.py

@@ -16,28 +16,45 @@ from models.world_model import WorldModel
 @hydra.main(config_path="../config", config_name="trainer")
 def main(cfg: DictConfig):
     device = torch.device(cfg.common.device)
-    assert cfg.mode in ('world_model', 'episode_replay', 'agent')
+    assert cfg.mode in ('episode_replay', 'agent_in_env', 'agent_in_world_model', 'play_in_world_model')
 
-    if cfg.mode in ['world_model', 'agent']:
-        env_fn = partial(instantiate, config=cfg.env.test)
-        test_env = SingleProcessEnv(env_fn)
-        tokenizer = instantiate(cfg.tokenizer)
-        world_model = WorldModel(obs_vocab_size=tokenizer.vocab_size, act_vocab_size=test_env.num_actions, config=instantiate(cfg.world_model))
-        if cfg.mode == 'world_model':
-            env = WorldModelEnv(tokenizer=tokenizer, world_model=world_model, pretrained_agent_path=Path('checkpoints/last.pt'), device=device, env=env_fn())
-            keymap = cfg.env.keymap
-        else:
-            actor_critic = ActorCritic(**cfg.actor_critic, act_vocab_size=test_env.num_actions)
-            agent = Agent(tokenizer, world_model, actor_critic).to(device)
-            agent.load(Path('checkpoints/last.pt'), device)
-            env = AgentEnv(agent, test_env, cfg.env.keymap)
-            keymap = 'empty'
+    env_fn = partial(instantiate, config=cfg.env.test)
+    test_env = SingleProcessEnv(env_fn)
 
+    if cfg.mode.startswith('agent_in_'):
+        h, w, _ = test_env.env.unwrapped.observation_space.shape
     else:
+        h, w = 64, 64
+    multiplier = 800 // h
+    size = [h * multiplier, w * multiplier]
+    
+    if cfg.mode == 'episode_replay':
         env = EpisodeReplayEnv(replay_keymap_name=cfg.env.keymap, episode_dir=Path('media/episodes'))
         keymap = 'episode_replay'
 
-    game = Game(env, keymap_name=keymap, size=(600, 1200 if cfg.mode == 'agent' else 600), fps=cfg.fps, verbose=bool(cfg.header))
+    else:
+        tokenizer = instantiate(cfg.tokenizer)
+        world_model = WorldModel(obs_vocab_size=tokenizer.vocab_size, act_vocab_size=test_env.num_actions, config=instantiate(cfg.world_model))
+        actor_critic = ActorCritic(**cfg.actor_critic, act_vocab_size=test_env.num_actions)
+        agent = Agent(tokenizer, world_model, actor_critic).to(device)
+        agent.load(Path('checkpoints/last.pt'), device)        
+
+        if cfg.mode == 'play_in_world_model':
+            env = WorldModelEnv(tokenizer=agent.tokenizer, world_model=agent.world_model, device=device, env=env_fn())
+            keymap = cfg.env.keymap
+        
+        elif cfg.mode == 'agent_in_env':
+            env = AgentEnv(agent, test_env, cfg.env.keymap, do_reconstruction=cfg.reconstruction)
+            keymap = 'empty'
+            if cfg.reconstruction:
+                size[1] *= 3
+
+        elif cfg.mode == 'agent_in_world_model':
+            wm_env = WorldModelEnv(tokenizer=agent.tokenizer, world_model=agent.world_model, device=device, env=env_fn())
+            env = AgentEnv(agent, wm_env, cfg.env.keymap, do_reconstruction=False)
+            keymap = 'empty'
+
+    game = Game(env, keymap_name=keymap, size=size, fps=cfg.fps, verbose=bool(cfg.header), record_mode=bool(cfg.save_mode))
     game.run()
 
 

src/utils.py

@@ -1,4 +1,5 @@
 from collections import OrderedDict
+import cv2
 from pathlib import Path
 import random
 import shutil
@@ -140,3 +141,14 @@ class RandomHeuristic:
         assert obs.ndim == 4  # (N, H, W, C)
         n = obs.size(0)
         return torch.randint(low=0, high=self.num_actions, size=(n,))
+
+
+def make_video(fname, fps, frames):
+    assert frames.ndim == 4 # (t, h, w, c)
+    t, h, w, c = frames.shape
+    assert c == 3
+
+    video = cv2.VideoWriter(str(fname), cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+    for frame in frames:
+        video.write(frame[:, :, ::-1])
+    video.release()