Spaces:
Running
Running
| import argparse | |
| import json | |
| import time | |
| import numpy as np | |
| import torch | |
| from pathlib import Path | |
| from utils.babyai_utils.baby_agent import load_agent | |
| from utils.env import make_env | |
| from utils.other import seed | |
| from utils.storage import get_model_dir | |
| from utils.storage import get_status | |
| from models import * | |
| import subprocess | |
| # Parse arguments | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--model", required=True, | |
| help="name of the trained model (REQUIRED)") | |
| parser.add_argument("--seed", type=int, default=0, | |
| help="random seed (default: 0)") | |
| parser.add_argument("--max-steps", type=int, default=None, | |
| help="max num of steps") | |
| parser.add_argument("--shift", type=int, default=0, | |
| help="number of times the environment is reset at the beginning (default: 0)") | |
| parser.add_argument("--argmax", action="store_true", default=False, | |
| help="select the action with highest probability (default: False)") | |
| parser.add_argument("--pause", type=float, default=0.5, | |
| help="pause duration between two consequent actions of the agent (default: 0.5)") | |
| parser.add_argument("--env-name", type=str, default=None, required=True, | |
| help="env name") | |
| parser.add_argument("--gif", type=str, default=None, | |
| help="store output as gif with the given filename", required=True) | |
| parser.add_argument("--episodes", type=int, default=10, | |
| help="number of episodes to visualize") | |
| args = parser.parse_args() | |
| # Set seed for all randomness sources | |
| seed(args.seed) | |
| save = args.gif | |
| if save: | |
| savename = args.gif | |
| if savename == "model_id": | |
| savename = args.model.replace('storage/', '') | |
| savename = savename.replace('/','_') | |
| savename += '_{}'.format(args.seed) | |
| # Set device | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"Device: {device}\n") | |
| # Load environment | |
| if str(args.model).startswith("./storage/"): | |
| args.model = args.model.replace("./storage/", "") | |
| if str(args.model).startswith("storage/"): | |
| args.model = args.model.replace("storage/", "") | |
| with open(Path("./storage") / args.model / "config.json") as f: | |
| conf = json.load(f) | |
| if args.env_name is None: | |
| # load env_args from status | |
| env_args = {} | |
| if not "env_args" in conf.keys(): | |
| env_args = get_status(get_model_dir(args.model), None)['env_args'] | |
| else: | |
| env_args = conf["env_args"] | |
| env = make_env(args.env_name, args.seed, env_args=env_args) | |
| else: | |
| env_name = args.env_name | |
| env = make_env(args.env_name, args.seed) | |
| for _ in range(args.shift): | |
| env.reset() | |
| print("Environment loaded\n") | |
| # Define agent | |
| model_dir = get_model_dir(args.model) | |
| num_frames = None | |
| agent = load_agent(env, model_dir, args.argmax, num_frames) | |
| print("Agent loaded\n") | |
| # Run the agent | |
| if save: | |
| from imageio import mimsave | |
| old_frames = [] | |
| frames = [] | |
| # Create a window to view the environment | |
| env.render(mode='human') | |
| def plt_2_rgb(env): | |
| data = np.frombuffer(env.window.fig.canvas.tostring_rgb(), dtype=np.uint8) | |
| data = data.reshape(env.window.fig.canvas.get_width_height()[::-1] + (3,)) | |
| return data | |
| for episode in range(args.episodes): | |
| print("episode:", episode) | |
| obs = env.reset() | |
| env.render(mode='human') | |
| if save: | |
| frames.append(plt_2_rgb(env)) | |
| i = 0 | |
| while True: | |
| i += 1 | |
| action = agent.get_action(obs) | |
| obs, reward, done, _ = env.step(action) | |
| agent.analyze_feedback(reward, done) | |
| env.render(mode='human') | |
| if save: | |
| img = plt_2_rgb(env) | |
| frames.append(img) | |
| if done: | |
| # quadruple last frame to pause between episodes | |
| for i in range(3): | |
| same_img = np.copy(img) | |
| # toggle a pixel between frames to avoid cropping when going from gif to mp4 | |
| same_img[0,0,2] = 0 if (i % 2) == 0 else 255 | |
| frames.append(same_img) | |
| if done or env.window.closed: | |
| break | |
| if args.max_steps is not None: | |
| if i > args.max_steps: | |
| break | |
| if env.window.closed: | |
| break | |
| if save: | |
| # from IPython import embed; embed() | |
| print(f"Saving to {savename} ", end="") | |
| mimsave(savename + '.gif', frames, duration=args.pause) | |
| # Reduce gif size | |
| # bashCommand = "gifsicle -O3 --colors 32 -o {}.gif {}.gif".format(savename, savename) | |
| # process = subprocess.run(bashCommand.split(), stdout=subprocess.PIPE) | |
| print("Done.") | |