def record(log_path,
env,
horizon,
kwargs,
modes=['real']):
_sanity_check(kwargs, modes)
init = env.reset()
Os = {}
total_costs = {}
for mode in modes:
Os[mode] = []
total_costs[mode] = []
output_path = osp.join(log_path, "%s.mp4" % mode)
encoder = ImageEncoder(output_path=output_path,
frame_shape=frame_size + (3,),
frames_per_sec=60)
print("Generating %s" % output_path)
obs = init
Os[mode].append(obs)
inner_env = _get_inner_env(env)
if mode == 'model':
inner_env.reset(obs)
image = inner_env.render(mode='rgb_array')
total_cost = 0.0
total_costs[mode].append(total_cost)
for t in range(horizon):
compressed_image = to_img(image, frame_size=frame_size)
# cv2.imshow('frame{}'.format(t), compressed_image)
# cv2.waitKey(10)
encoder.capture_frame(compressed_image)
action = _get_action(kwargs, obs)
action = np.clip(action, *env.action_space.bounds)
next_obs, reward, done, info = _step(kwargs, env, obs, action, mode)
total_cost -= reward
obs = next_obs
Os[mode].append(obs)
if mode == 'model':
inner_env.reset(next_obs)
image = inner_env.render(mode='rgb_array')
# if done:
# break
total_costs[mode].append(total_cost)
print("%s cost: %f" % (mode, total_cost))
encoder.close()
if len(Os) == 2:
_analyze_trajectories(Os, total_costs, log_path)
def animate(self, act_fn, nsteps, **kwargs):
"""act_fn could be a list of functions for each agent in the environemnt that we can control"""
if not isinstance(act_fn, list):
act_fn = [act_fn for _ in range(len(self.agents))]
assert len(act_fn) == len(self.agents)
encoder = None
vid_loc = kwargs.pop('vid', None)
obs = self.reset()
frame = self.render(**kwargs)
if vid_loc:
fps = kwargs.pop('fps', 30)
encoder = ImageEncoder(vid_loc, frame.shape, fps)
try:
encoder.capture_frame(frame)
except error.InvalidFrame as e:
print('Invalid video frame, {}'.format(e))
rew = np.zeros((len(self.agents)))
traj_info_list = []
for step in range(nsteps):
a = list(map(lambda afn, o: afn(o), act_fn, obs))
obs, r, done, info = self.step(a)
rew += r
if info:
traj_info_list.append(info)
frame = self.render(**kwargs)
if vid_loc:
try:
encoder.capture_frame(frame)
except error.InvalidFrame as e:
print('Invalid video frame, {}'.format(e))
if done:
break
traj_info = stack_dict_list(traj_info_list)
return rew, traj_info
def animate(self, act_fn, nsteps, **kwargs):
"""act_fn could be a list of functions for each agent in the environemnt that we can control"""
if not isinstance(act_fn, list):
act_fn = [act_fn for _ in range(len(self.agents))]
assert len(act_fn) == len(self.agents)
encoder = None
vid_loc = kwargs.pop('vid', None)
obs = self.reset()
frame = self.render(**kwargs)
if vid_loc:
fps = kwargs.pop('fps', 30)
encoder = ImageEncoder(vid_loc, frame.shape, fps)
try:
encoder.capture_frame(frame)
except error.InvalidFrame as e:
print('Invalid video frame, {}'.format(e))
rew = np.zeros((len(self.agents)))
traj_info_list = []
for step in range(nsteps):
a = list(map(lambda afn, o: afn(o), act_fn, obs))
obs, r, done, info = self.step(a)
rew += r
if info:
traj_info_list.append(info)
frame = self.render(**kwargs)
if vid_loc:
try:
encoder.capture_frame(frame)
except error.InvalidFrame as e:
print('Invalid video frame, {}'.format(e))
if done:
break
traj_info = stack_dict_list(traj_info_list)
return rew, traj_info
policy = data["policy"]
env = data["env"]
# env = SwimmerEnv()
for idx in range(7, 8):
encoder = ImageEncoder(output_path=osp.join(
output_path, '%d_goalGAN_maze.mp4' % idx),
frame_shape=frame_size + (3, ),
frames_per_sec=15)
for i in range(6):
obs = env.reset()
print("Generating %d_goalGAN_maze.mp4" % idx)
image = env.render(mode='rgb_array')
policy.reset()
for t in range(500):
compressed_image = to_img(image, frame_size=frame_size)
# cv2.imshow('frame{}'.format(t), compressed_image)
cv2.waitKey(10)
encoder.capture_frame(compressed_image)
action, _ = policy.get_action(obs)
next_obs, reward, done, info = env.step(action)
obs = next_obs
image = env.render(mode='rgb_array')
if done:
break
encoder.close()
path = rollout(env, policy, 30, True, 999, None, 'rgb_array', viewer_settings)
for p in range(n_paths):
path = rollout(
env=env,
policy=policy,
path_length=path_length,
speedup=9999,
render=True,
render_mode='rgb_array',
viewer_kwargs=viewer_settings
)
ims.append(path['ims'])
ims = np.concatenate(ims, axis=0)
video_file = osp.join(output_path, name + '.mp4')
encoder = ImageEncoder(
output_path=video_file,
frame_shape=frame_size + (3,),
frames_per_sec=20
)
for im in ims:
encoder.capture_frame(im)
encoder.close()
tf.reset_default_graph()