<jupyter_start><jupyter_text>Train policyThe training of the policy will take around 1h on GPU.<jupyter_code>iteration_num=300
hparams = trainer_lib.create_hparams("ppo_atari_base", "epochs_num={}".format(iteration_num+1))
ppo_dir = tempfile.mkdtemp(dir=data_dir, prefix="ppo_")
rl_trainer_lib.train(hparams, "stacked_pong", ppo_dir)
agent_policy_path = os.path.join(ppo_dir, "model{}.ckpt.index".format(iteration_num))[:-6]<jupyter_output><empty_output><jupyter_text>Generate and review frames from policy<jupyter_code>sys.argv = [sys.argv[0], "--agent_policy_path", agent_policy_path]
tf.reset_default_graph()
# This step is also time consuming - takes around 30 minutes.
gym_problem = problems.problem("gym_pong_trajectories_from_policy")
gym_problem.generate_data(data_dir, tmp_dir)
dataset = gym_problem.dataset(Modes.TRAIN, data_dir)
iterator = dataset.make_one_shot_iterator()
next_element = iterator.get_next()


fig=plt.figure(figsize=(20, 80))
columns = 10
rows = 40

with tf.Session() as sess:
    for inx in range(100):
        value = sess.run(next_element)
        for i in range(10):  # skipping surplus frames.
            value = sess.run(next_element)
        fig.add_subplot(rows, columns, inx+1)        
        image = value["inputs"].reshape([210,160,3])
        plt.imshow(image[:, :, 0].astype(np.float32), cmap=plt.get_cmap('gray'))
plt.show()<jupyter_output>INFO:tensorflow:Reading data files from /home/blazej.osinski/t2t_rl_data/data/gym_pong_trajectories_from_policy-train*
{'fraction_non_alphanumeric': 0.1284953395472703, 'fraction_numerical': 0.022636484687083888, 'mean_word_length': 5.831818181818182, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '221259', 'license': 'Apache-2.0', 'n_tokens_mistral': 576, 'n_tokens_neox': 540, 'n_words': 92, 'path': 'tensor2tensor/notebooks/hello_t2t-rl.ipynb', 'repo_name': 'repoloper/tensor2tensor', 'size': '1502'}