<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Train policy
The training of the policy will take around 1h on GPU.
Step1:  Generate and review frames from policy
<ASSISTANT_TASK:>
Python 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]
Explanation: Train policy
The training of the policy will take around 1h on GPU.
End of explanation
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()
Explanation: Generate and review frames from policy
End of explanation
<END_TASK>