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

README.md

@@ -100,9 +100,13 @@ outputs/YYYY-MM-DD/hh-mm-ss/
   - `eval.py`: Launch `python ./scripts/eval.py` to evaluate the run.
   - `resume.sh`: Launch `./scripts/resume.sh` to resume a training that crashed.
   - `play.sh`: Tool to visualize some interesting aspects of the run.
-    - Launch `./scripts/play.sh -a` to watch the agent play live in the environment. The left panel displays the original environment, and the right panel shows what the agent actually sees through its discrete autoencoder.
+    - Launch `./scripts/play.sh` to watch the agent play live in the environment. If you add the flag `-r`, the left panel displays the original frame, the center panel displays the same frame downscaled to the input resolution of the discrete autoencoder, and the right panel shows the output of the autoencoder (what the agent actually sees).
     - Launch `./scripts/play.sh -w` to unroll live trajectories with your keyboard inputs (i.e. to play in the world model). Note that for faster interaction, the memory of the Transformer is flushed every 20 frames.
-    - Launch `./scripts/play.sh` to visualize the episodes contained in `media/episodes`.
+    - Launch `./scripts/play.sh -a` to watch the agent play live in the world model. Note that for faster interaction, the memory of the Transformer is flushed every 20 frames.
+    - Launch `./scripts/play.sh -e` to visualize the episodes contained in `media/episodes`.
+    - Add the flag `-h` to display a header with additional information.
+    - Press '`,`' to start and stop recording. The corresponding segment is saved in `media/recordings` in mp4 and numpy formats.
+    - Add the flag `-s` to enter 'save mode', where the user is prompted to save trajectories upon completion.
 
 ## Results notebook
 

scripts/play.sh

@@ -1,8 +1,10 @@
 #!/usr/bin/bash
 
 fps=15
-header=1
-mode="episode_replay"
+header=0
+reconstruction=0
+save_mode=0
+mode="agent_in_env"
 
 while [ "$1" != "" ]; do
     case $1 in
@@ -10,14 +12,23 @@ while [ "$1" != "" ]; do
             shift
             fps=$1
             ;;
-        -h | --no-header )
-            header=0
+        -h | --header )
+            header=1
+            ;;
+        -r | --reconstruction )
+            reconstruction=1
+            ;;
+        -s | --save-mode )
+            save_mode=1
+            ;;
+        -a | --agent-world-model )
+            mode="agent_in_world_model"
+            ;;
+        -e | --episode )
+            mode="episode_replay"
             ;;
         -w | --world-model )
-            mode="world_model"
-            ;;
-        -a | --agent )
-            mode="agent"
+            mode="play_in_world_model"
             ;;
         * )
             echo Invalid usage : $1
@@ -26,4 +37,4 @@ while [ "$1" != "" ]; do
     shift
 done
 
-python src/play.py hydra.run.dir=. hydra.output_subdir=null +mode="${mode}" +fps="${fps}" +header="${header}"
+python src/play.py hydra.run.dir=. hydra.output_subdir=null +mode="${mode}" +fps="${fps}" +header="${header}" +reconstruction="${reconstruction}" +save_mode="${save_mode}"

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()