def _setup(self, variant):
# Set the current working directory such that the local mode
# logs into the correct place. This would not be needed on
# local/cluster mode.
if ray.worker._mode() == ray.worker.LOCAL_MODE:
os.chdir(os.getcwd())
set_seed(variant['run_params']['seed'])
if variant['run_params'].get('run_eagerly', False):
tf.config.experimental_run_functions_eagerly(True)
self._variant = variant
set_gpu_memory_growth(True)
self.train_generator = None
self._built = False
render_kwargs,
video_save_path=None,
evaluation_environment_params=None):
checkpoint_path = os.path.abspath(checkpoint_path.rstrip('/'))
variant, progress, metadata = load_variant_progress_metadata(
checkpoint_path)
environment = load_environment(variant)
policy = load_policy(checkpoint_path, variant, environment)
render_kwargs = {**DEFAULT_RENDER_KWARGS, **render_kwargs}
paths = rollouts(num_rollouts,
environment,
policy,
path_length=max_path_length,
render_kwargs=render_kwargs)
if video_save_path and render_kwargs.get('mode') == 'rgb_array':
fps = 1 // getattr(environment, 'dt', 1 / 30)
for i, path in enumerate(paths):
video_save_dir = os.path.expanduser('/tmp/simulate_policy/')
video_save_path = os.path.join(video_save_dir, f'episode_{i}.mp4')
save_video(path['images'], video_save_path, fps=fps)
return paths
if __name__ == '__main__':
set_gpu_memory_growth(True)
args = parse_args()
simulate_policy(**vars(args))
