PyTorch 1.x Reinforcement Learning Cookbook

There's more...

In fact, the optimal policy was obtained after around 50 episodes. We can plot the length of each episode over time to verify this. The total reward obtained in each episode over time is also an option.

1. We define two lists to store the length and total reward for each episode, respectively:

>>> length_episode = [0] * n_episode>>> total_reward_episode = [0] * n_episode

2. We keep track of the length and total reward for each episode during learning. The following is the updated version of q_learning:

>>> def q_learning(env, gamma, n_episode, alpha): ...     n_action = env.action_space.n ...     Q = defaultdict(lambda: torch.zeros(n_action)) ...     for episode in range(n_episode): ...         state = env.reset() ...         is_done = False ... while ...