human_feedback.py 文件源码

def online_receive_frame(viewer_state, episode):
    logger.info('Waiting for message...')
    if not viewer_state.conn.poll(1000.0):
        raise EnvironmentError("Failed to receive message!")
    msg = viewer_state.conn.recv()
    if isinstance(msg, frame_module.Frame):
        viewer_state.current_frame = msg
        episode.frames.append(viewer_state.current_frame)
    elif msg['msg'] == 'init':
        viewer_state.prev_frame = msg
        viewer_state.action_set = msg['action_set']
        viewer_state.episode_num_in_session += 1
        print('Initial message received')
        print('... got action_set: {}'.format(viewer_state.action_set))
        print('... episode_num: {}'.format(viewer_state.episode_num_in_session))
        viewer_state.env_id = msg['env_id']
        viewer_state.skip_frame = True
    elif msg['msg'] == 'close':
        viewer_state.conn.close()
        viewer_state.close = True
    elif msg['msg'] == 'done':
        viewer_state.frame_index = 0
        proceed = input('Proceed to next episode?')

        if proceed != 'y':
            viewer_state.EXIT = True

        while True:
            k = cv2.waitKey(viewer_state.delay) & 0xFF
            if k == 255: break
        viewer_state.skip_episode = True
    else:
        print('Unknown message received: {}'.format(msg))
        viewer_state.skip_frame = True

    return viewer_state, episode