Beispiel #1
0
    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
Beispiel #2
0
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)
Beispiel #3
0
    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
Beispiel #4
0
    )
    output_path = osp.join(config.PROJECT_PATH, "data/video/goalGAN_maze")
    console.mkdir_p(output_path)

    # import pdb; pdb.set_trace()
    data = joblib.load(pkl_file)

    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
Beispiel #5
0
    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()