def main(): args = arg_parser().parse_args() print('Creating environments...') env = create_env(args.env, 1, 1, args.fps, args.max_timesteps) env = ObsInInfo(env) env = wrap_env(env) try: print('Creating session...') with tf.Session() as sess: print('Creating model graph...') model = IMPALAModel(sess, *gym_spaces(env)) print('Initializing model variables...') sess.run(tf.global_variables_initializer()) load_vars(sess, args.save_path) print('Gathering episode...') def run_episode(): states = model.start_state(1) env.reset_start() obses = env.reset_wait() while True: outputs = model.step(obses, states) env.step_start(outputs['actions']) obses, _, dones, infos = env.step_wait() states = outputs['states'] yield pad_dims(infos[0]['old_obs']) if dones[0]: return spec = muniverse.spec_for_name(args.env) export_video(args.path, padded_dim(spec['Width']), padded_dim(spec['Height']), args.fps, run_episode()) finally: env.close()
def main(): def image_fn(): env = StateEnv(create_single_env(random.randrange(15, 20), clear=False)) model = ACModel() tail_model = ACModel() model.load_state_dict(torch.load('save.pkl', map_location='cpu')) tail_model.load_state_dict(torch.load('save_tail.pkl', map_location='cpu')) model.to(torch.device('cuda')) tail_model.to(torch.device('cuda')) state, obs = env.reset() floor = 0 while True: output = (model if floor < 10 else tail_model).step(np.array([state]), np.array([obs])) (state, obs), rew, done, info = env.step(output['actions'][0]) floor = info['current_floor'] yield big_obs(obs[..., -3:], info) if done: break env.close() export_video('export.mp4', 168, 168, 10, image_fn())
import json import os from anyrl.utils.ffmpeg import export_video import numpy as np from obstacle_tower_env import ObstacleTowerEnv from obs_tower2.util import big_obs with open('stuck_box.json', 'r') as in_file: data = json.load(in_file) env = ObstacleTowerEnv(os.environ['OBS_TOWER_PATH'], worker_id=1) env.seed(56) env.reset() def f(): for i, act in enumerate(data): obs, _, _, info = env.step(act) if i > 5275: yield big_obs(obs, info) export_video('stuck_box.mp4', 168, 168, 10, f())
import json import os from anyrl.utils.ffmpeg import export_video import numpy as np from obstacle_tower_env import ObstacleTowerEnv from obs_tower2.util import big_obs with open('hang.json', 'r') as in_file: data = json.load(in_file) env = ObstacleTowerEnv(os.environ['OBS_TOWER_PATH'], worker_id=4) env.reset() def f(): env.seed(data['seed']) env.floor(data['floor']) env.reset() for act in data['actions'][:-1]: obs, _, _, info = env.step(act) yield big_obs(obs, info) export_video('hang.mp4', 168, 168, 10, f())