Code release.

scripts/eval.py

@@ -0,0 +1,57 @@
+#! python
+
+import argparse
+from datetime import datetime
+from pathlib import Path
+import subprocess
+
+from omegaconf import OmegaConf
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-n', '--num-episodes', type=int, default=100, help='Number of evaluation episodes to collect.')
+    parser.add_argument('-p', '--num-envs', type=int, default=25, help='Number of environments used to collect the evaluation episodes.')
+    args = parser.parse_args()
+
+    path_to_config = Path('config') / 'trainer.yaml'
+    cfg = OmegaConf.load(path_to_config)
+
+    path_to_checkpoint = Path('checkpoints') / 'last.pt'
+    assert path_to_checkpoint.is_file()
+
+    cmd = f'python src/main.py hydra.run.dir=eval_outputs/{datetime.now().strftime("%Y-%m-%d/%H-%M-%S")} '
+
+    cmd += 'wandb.mode=online '
+    cmd += f'wandb.name=eval-{cfg.wandb.name} '
+    cmd += f'wandb.group=eval-{cfg.wandb.group} '
+
+    cmd += f'initialization.path_to_checkpoint={str(path_to_checkpoint.absolute())} '
+    cmd += 'initialization.load_tokenizer=True '
+    cmd += 'initialization.load_world_model=False '
+    cmd += 'initialization.load_actor_critic=True '
+
+    cmd += 'common.epochs=1 '
+    cmd += 'common.device=cuda:0 '
+    cmd += 'common.do_checkpoint=False '
+    cmd += 'common.seed=0 '
+
+    cmd += 'collection.test.num_episodes_to_save=0 '
+    cmd += f'collection.test.num_envs={args.num_envs} '
+    cmd += f'collection.test.config.num_episodes={args.num_episodes} '
+
+    cmd += 'training.should=False '
+
+    # Turn on data collection only
+    cmd += 'evaluation.should=True '
+    cmd += 'evaluation.every=1 '
+    cmd += 'evaluation.tokenizer.start_after_epochs=1 '
+    cmd += 'evaluation.tokenizer.save_reconstructions=False '
+    cmd += 'evaluation.world_model.start_after_epochs=1 '
+    cmd += 'evaluation.actor_critic.start_after_epochs=1 '
+
+    subprocess.run(cmd, shell=True, check=True)
+
+
+if __name__ == '__main__':
+    main()

scripts/import_run_folder.py

@@ -0,0 +1,60 @@
+#! /usr/bin/python3
+
+import argparse
+from pathlib import Path
+import shutil
+import subprocess
+import yaml
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('src_folder', type=str, help='Folder to import. Must be formatted as: USERNAME@HOSTNAME:PATH_TO_OUTPUTS_FOLDER/DATE/TIME')
+    parser.add_argument('-k', '--from-key', action='store_true')
+    parser.add_argument('-v', '--verbose', action='store_true')
+    args = parser.parse_args()
+
+    if args.from_key:
+        key_file = Path('path_runs.yaml')
+        assert key_file.is_file()
+        with key_file.open('r') as f:
+            runs = yaml.load(f, Loader=yaml.loader.SafeLoader)
+        src_folder = runs[args.src_folder]
+    else:
+        src_folder = args.src_folder
+
+    day, time = src_folder.split('/')[-2:]
+    host = src_folder.split('@')[1].split(':')[0]
+    dst_folder = Path(host) / day / time
+    print(dst_folder)
+    if dst_folder.is_dir():
+        if input(f'{dst_folder} exists, remove it ? [Y/n] ').lower() != 'n':
+            shutil.rmtree(dst_folder)
+        else:
+            print('Bye.')
+            return
+
+    dst_folder.mkdir(exist_ok=False, parents=True)
+
+    # Make symbolic link from key to folder
+    if args.from_key:
+        subprocess.run(f'ln -s {str(dst_folder)} {args.src_folder}', shell=True, check=True)
+
+    download_media = args.verbose and input('Download media/ ? [y/N] ').lower() == 'y'
+    download_last = not args.verbose or input('Download checkpoints/last.pt ? [Y/n] ').lower() != 'n'
+
+    folders = ['src', 'config', 'scripts']
+    if download_media:
+        folders.append('media')
+
+    for folder in folders:
+        subprocess.run(f'scp -r {src_folder}/{folder} {dst_folder}', shell=True, check=True)
+
+    if download_last:
+        checkpoint_folder = dst_folder / 'checkpoints'
+        checkpoint_folder.mkdir(exist_ok=False, parents=False)
+        subprocess.run(f'scp {src_folder}/checkpoints/last.pt {checkpoint_folder}', shell=True, check=True)
+
+
+if __name__ == '__main__':
+    main()

scripts/play.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/bash
+
+fps=15
+header=1
+mode="episode_replay"
+
+while [ "$1" != "" ]; do
+    case $1 in
+        -f | --fps )
+            shift
+            fps=$1
+            ;;
+        -h | --no-header )
+            header=0
+            ;;
+        -w | --world-model )
+            mode="world_model"
+            ;;
+        -a | --agent )
+            mode="agent"
+            ;;
+        * )
+            echo Invalid usage : $1
+            exit 1
+    esac
+    shift
+done
+
+python src/play.py hydra.run.dir=. hydra.output_subdir=null +mode="${mode}" +fps="${fps}" +header="${header}"

scripts/resume.sh

@@ -0,0 +1 @@
+python src/main.py common.resume=True hydra.output_subdir=null hydra.run.dir=.