Safe lab preview

Cart Pole-RL-Py Torch

This is a sanitized, read-only preview. Nothing executes in this page.

Read-only

Notebook preview

Cart Pole-RL-Py Torch

Download notebook

Markdown cell[1]

> **Integrated runtime note:** This preview uses a small deterministic, rights-safe offline fixture for reproducible learning. Full-scale results require the lesson's documented dataset or model in an approved external environment. # Training RL to do Cartpole Balancing This notebook is part of AI for Beginners. It has been inspired by [official PyTorch tutorial](https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html) and [this Cartpole PyTorch implementation](https://github.com/yc930401/Actor-Critic-pytorch). In this example, we will use RL to train a model to balance a pole on a cart that can move left and right on horizontal scale. We will use [Gymnasium](https://gymnasium.farama.org/) environment to simulate the pole. > **Note**: This edition uses the current Gymnasium API. In hosted environments, prefer `rgb_array` output instead of relying on a local graphical window. We will start by making sure Gym is installed:

Code cell[2]

import sys
# Dependencies are provisioned by this course edition's pinned runtime profile.

Markdown cell[3]

Now let's create the CartPole environment and see how to operate on it. An environment has the following properties: * **Action space** is the set of possible actions that we can perform at each step of the simulation * **Observation space** is the space of observations that we can make

Code cell[4]

import gymnasium as gym

env = gym.make("CartPole-v1", render_mode="rgb_array")

print(f"Action space: {env.action_space}")
print(f"Observation space: {env.observation_space}")

Markdown cell[5]

Let's see how the simulation works. The following loop runs the simulation, until `env.step` does not return the termination flag `done`. We will randomly chose actions using `env.action_space.sample`, which means the experiment will probably fail very fast (CartPole environment terminates when the speed of CartPole, its position or angle are outside certain limits). > Simulation will open in the new window. You can run the code several times and see how it behaves.

Code cell[6]

env.reset()

done = False
total_reward = 0
while not done:
   env.render()
   obs, rew, terminated, truncated, info = env.step(env.action_space.sample())
   done = terminated or truncated
   total_reward += rew
   print(f"{obs} -> {rew}")
print(f"Total reward: {total_reward}")

Markdown cell[7]

Youn can notice that observations contain 4 numbers. They are: - Position of cart - Velocity of cart - Angle of pole - Rotation rate of pole `rew` is the reward we receive at each step. You can see that in CartPole environment you are rewarded 1 point for each simulation step, and the goal is to maximize total reward, i.e. the time CartPole is able to balance without falling. During reinforcement learning, our goal is to train a **policy** $\pi$, that for each state $s$ will tell us which action $a$ to take, so essentially $a = \pi(s)$. If you want probabilistic solution, you can think of policy as returning a set of probabilities for each action, i.e. $\pi(a|s)$ would mean a probability that we should take action $a$ at state $s$. ## Policy Gradient Method In simplest RL algorithm, called **Policy Gradient**, we will train a neural network to predict the next action.

Code cell[8]

import numpy as np
import matplotlib.pyplot as plt
import torch

num_inputs = 4
num_actions = 2

model = torch.nn.Sequential(
    torch.nn.Linear(num_inputs, 128, bias=False, dtype=torch.float32),
    torch.nn.ReLU(),
    torch.nn.Linear(128, num_actions, bias = False, dtype=torch.float32),
    torch.nn.Softmax(dim=1)
)

np.random.seed(2026)

Markdown cell[9]

We will train the network by running many experiments, and updating our network after each run. Let's define a function that will run the experiment and return the results (so-called **trace**) - all states, actions (and their recommended probabilities), and rewards:

Code cell[10]

def run_episode(max_steps_per_episode = 10000,render=False):
    states, actions, probs, rewards = [],[],[],[]
    state, _ = env.reset()
    for _ in range(max_steps_per_episode):
        if render:
            env.render()
        action_probs = model(torch.from_numpy(np.expand_dims(state,0)))[0]
        action = np.random.choice(num_actions, p=np.squeeze(action_probs.detach().numpy()))
        nstate, reward, terminated, truncated, info = env.step(action)
        done = terminated or truncated
        if done:
            break
        states.append(state)
        actions.append(action)
        probs.append(action_probs.detach().numpy())
        rewards.append(reward)
        state = nstate
    return np.vstack(states), np.vstack(actions), np.vstack(probs), np.vstack(rewards)

Markdown cell[11]

You can run one episode with untrained network and observe that total reward (AKA length of episode) is very low:

Code cell[12]

s, a, p, r = run_episode()
print(f"Total reward: {np.sum(r)}")

Markdown cell[13]

One of the tricky aspects of policy gradient algorithm is to use **discounted rewards**. The idea is that we compute the vector of total rewards at each step of the game, and during this process we discount the early rewards using some coefficient $gamma$. We also normalize the resulting vector, because we will use it as weight to affect our training:

Code cell[14]

eps = 0.0001

def discounted_rewards(rewards,gamma=0.99,normalize=True):
    ret = []
    s = 0
    for r in rewards[::-1]:
        s = r + gamma * s
        ret.insert(0, s)
    if normalize:
        ret = (ret-np.mean(ret))/(np.std(ret)+eps)
    return ret

Markdown cell[15]

Now let's do the actual training! We will run 300 episodes, and at each episode we will do the following: 1. Run the experiment and collect the trace 1. Calculate the difference (`gradients`) between the actions taken, and by predicted probabilities. The less the difference is, the more we are sure that we have taken the right action. 1. Calculate discounted rewards and multiply gradients by discounted rewards - that will make sure that steps with higher rewards will make more effect on the final result than lower-rewarded ones 1. Expected target actions for our neural network would be partly taken from the predicted probabilities during the run, and partly from calculated gradients. We will use `alpha` parameter to determine to which extent gradients and rewards are taken into account - this is called *learning rate* of reinforcement algorithm. 1. Finally, we train our network on states and expected actions, and repeat the process

Code cell[16]

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

def train_on_batch(x, y):
    x = torch.from_numpy(x)
    y = torch.from_numpy(y)
    optimizer.zero_grad()
    predictions = model(x)
    loss = -torch.mean(torch.log(predictions) * y)
    loss.backward()
    optimizer.step()
    return loss

Code cell[17]

alpha = 1e-4

history = []
for epoch in range(300):
    states, actions, probs, rewards = run_episode()
    one_hot_actions = np.eye(2)[actions.T][0]
    gradients = one_hot_actions-probs
    dr = discounted_rewards(rewards)
    gradients *= dr
    target = alpha*np.vstack([gradients])+probs
    train_on_batch(states,target)
    history.append(np.sum(rewards))
    if epoch%100==0:
        print(f"{epoch} -> {np.sum(rewards)}")

plt.plot(history)

Markdown cell[18]

Now let's run the episode with rendering to see the result:

Code cell[19]

_ = run_episode(render=True)

Markdown cell[20]

Hopefully, you can see that pole can now balance pretty well! ## Actor-Critic Model Actor-Critic model is the further development of policy gradients, in which we build a neural network to learn both the policy and estimated rewards. The network will have two outputs (or you can view it as two separate networks): * **Actor** will recommend the action to take by giving us the state probability distribution, as in policy gradient model * **Critic** would estimate what the reward would be from those actions. It returns total estimated rewards in the future at the given state. Let's define such a model:

Code cell[21]

from itertools import count
import torch.nn.functional as F

Code cell[22]

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
env.close()
env = gym.make("CartPole-v1", render_mode="rgb_array")

state_size = env.observation_space.shape[0]
action_size = env.action_space.n
lr = 0.0001

class Actor(torch.nn.Module):
    def __init__(self, state_size, action_size):
        super(Actor, self).__init__()
        self.state_size = state_size
        self.action_size = action_size
        self.linear1 = torch.nn.Linear(self.state_size, 128)
        self.linear2 = torch.nn.Linear(128, 256)
        self.linear3 = torch.nn.Linear(256, self.action_size)

    def forward(self, state):
        output = F.relu(self.linear1(state))
        output = F.relu(self.linear2(output))
        output = self.linear3(output)
        distribution = torch.distributions.Categorical(F.softmax(output, dim=-1))
        return distribution


class Critic(torch.nn.Module):
    def __init__(self, state_size, action_size):
        super(Critic, self).__init__()
        self.state_size = state_size
        self.action_size = action_size
        self.linear1 = torch.nn.Linear(self.state_size, 128)
        self.linear2 = torch.nn.Linear(128, 256)
        self.linear3 = torch.nn.Linear(256, 1)

    def forward(self, state):
        output = F.relu(self.linear1(state))
        output = F.relu(self.linear2(output))
        value = self.linear3(output)
        return value

Markdown cell[23]

We would need to slightly modify our `discounted_rewards` and `run_episode` functions:

Code cell[24]

def discounted_rewards(next_value, rewards, masks, gamma=0.99):
    R = next_value
    returns = []
    for step in reversed(range(len(rewards))):
        R = rewards[step] + gamma * R * masks[step]
        returns.insert(0, R)
    return returns

def run_episode(actor, critic, n_iters):
    optimizerA = torch.optim.Adam(actor.parameters())
    optimizerC = torch.optim.Adam(critic.parameters())
    for iter in range(n_iters):
        state, _ = env.reset()
        log_probs = []
        values = []
        rewards = []
        masks = []
        entropy = 0
        env.reset()

        for i in count():
            env.render()
            state = torch.FloatTensor(state).to(device)
            dist, value = actor(state), critic(state)

            action = dist.sample()
            next_state, reward, terminated, truncated, _ = env.step(action.cpu().numpy())
            done = terminated or truncated
            log_prob = dist.log_prob(action).unsqueeze(0)
            entropy += dist.entropy().mean()

            log_probs.append(log_prob)
            values.append(value)
            rewards.append(torch.tensor([reward], dtype=torch.float, device=device))
            masks.append(torch.tensor([1-done], dtype=torch.float, device=device))

            state = next_state

            if done:
                print('Iteration: {}, Score: {}'.format(iter, i))
                break


        next_state = torch.FloatTensor(next_state).to(device)
        next_value = critic(next_state)
        returns = discounted_rewards(next_value, rewards, masks)

        log_probs = torch.cat(log_probs)
        returns = torch.cat(returns).detach()
        values = torch.cat(values)

        advantage = returns - values

        actor_loss = -(log_probs * advantage.detach()).mean()
        critic_loss = advantage.pow(2).mean()

        optimizerA.zero_grad()
        optimizerC.zero_grad()
        actor_loss.backward()
        critic_loss.backward()
        optimizerA.step()
        optimizerC.step()

Markdown cell[25]

Now we will run the main training loop. We will use manual network training process by computing proper loss functions and updating network parameters:

Code cell[26]


actor = Actor(state_size, action_size).to(device)
critic = Critic(state_size, action_size).to(device)
run_episode(actor, critic, n_iters=100)

Markdown cell[27]

Finally, let's close the environment.

Code cell[28]

env.close()

Markdown cell[29]

## Takeaway We have seen two RL algorithms in this demo: simple policy gradient, and more sophisticated actor-critic. You can see that those algorithms operate with abstract notions of state, action and reward - thus they can be applied to very different environments. Reinforcement learning allows us to learn the best strategy to solve the problem just by looking at the final reward. The fact that we do not need labelled datasets allows us to repeat simulations many times to optimize our models. However, there are still many challenges in RL, which you may learn if you decide to focus more on this interesting area of AI.

Outputs, execution counts, widgets, and active content were removed during import. Run notebooks only in an external environment you trust.

Record your practice

Optional self-reporting helps you remember what you practiced and never gates course completion.

Back to lesson

env.reset() done = False total_reward = 0 while not done: env.render() obs, rew, terminated, truncated, info = env.step(env.action_space.sample()) done = terminated or truncated total_reward += rew print(f"{obs} -> {rew}") print(f"Total reward: {total_reward}")

import numpy as np import matplotlib.pyplot as plt import torch num_inputs = 4 num_actions = 2 model = torch.nn.Sequential( torch.nn.Linear(num_inputs, 128, bias=False, dtype=torch.float32), torch.nn.ReLU(), torch.nn.Linear(128, num_actions, bias = False, dtype=torch.float32), torch.nn.Softmax(dim=1) ) np.random.seed(2026)

def run_episode(max_steps_per_episode = 10000,render=False): states, actions, probs, rewards = [],[],[],[] state, _ = env.reset() for _ in range(max_steps_per_episode): if render: env.render() action_probs = model(torch.from_numpy(np.expand_dims(state,0)))[0] action = np.random.choice(num_actions, p=np.squeeze(action_probs.detach().numpy())) nstate, reward, terminated, truncated, info = env.step(action) done = terminated or truncated if done: break states.append(state) actions.append(action) probs.append(action_probs.detach().numpy()) rewards.append(reward) state = nstate return np.vstack(states), np.vstack(actions), np.vstack(probs), np.vstack(rewards)

eps = 0.0001 def discounted_rewards(rewards,gamma=0.99,normalize=True): ret = [] s = 0 for r in rewards[::-1]: s = r + gamma * s ret.insert(0, s) if normalize: ret = (ret-np.mean(ret))/(np.std(ret)+eps) return ret

optimizer = torch.optim.Adam(model.parameters(), lr=0.01) def train_on_batch(x, y): x = torch.from_numpy(x) y = torch.from_numpy(y) optimizer.zero_grad() predictions = model(x) loss = -torch.mean(torch.log(predictions) * y) loss.backward() optimizer.step() return loss

alpha = 1e-4 history = [] for epoch in range(300): states, actions, probs, rewards = run_episode() one_hot_actions = np.eye(2)[actions.T][0] gradients = one_hot_actions-probs dr = discounted_rewards(rewards) gradients *= dr target = alpha*np.vstack([gradients])+probs train_on_batch(states,target) history.append(np.sum(rewards)) if epoch%100==0: print(f"{epoch} -> {np.sum(rewards)}") plt.plot(history)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") env.close() env = gym.make("CartPole-v1", render_mode="rgb_array") state_size = env.observation_space.shape[0] action_size = env.action_space.n lr = 0.0001 class Actor(torch.nn.Module): def __init__(self, state_size, action_size): super(Actor, self).__init__() self.state_size = state_size self.action_size = action_size self.linear1 = torch.nn.Linear(self.state_size, 128) self.linear2 = torch.nn.Linear(128, 256) self.linear3 = torch.nn.Linear(256, self.action_size) def forward(self, state): output = F.relu(self.linear1(state)) output = F.relu(self.linear2(output)) output = self.linear3(output) distribution = torch.distributions.Categorical(F.softmax(output, dim=-1)) return distribution class Critic(torch.nn.Module): def __init__(self, state_size, action_size): super(Critic, self).__init__() self.state_size = state_size self.action_size = action_size self.linear1 = torch.nn.Linear(self.state_size, 128) self.linear2 = torch.nn.Linear(128, 256) self.linear3 = torch.nn.Linear(256, 1) def forward(self, state): output = F.relu(self.linear1(state)) output = F.relu(self.linear2(output)) value = self.linear3(output) return value

def discounted_rewards(next_value, rewards, masks, gamma=0.99): R = next_value returns = [] for step in reversed(range(len(rewards))): R = rewards[step] + gamma * R * masks[step] returns.insert(0, R) return returns def run_episode(actor, critic, n_iters): optimizerA = torch.optim.Adam(actor.parameters()) optimizerC = torch.optim.Adam(critic.parameters()) for iter in range(n_iters): state, _ = env.reset() log_probs = [] values = [] rewards = [] masks = [] entropy = 0 env.reset() for i in count(): env.render() state = torch.FloatTensor(state).to(device) dist, value = actor(state), critic(state) action = dist.sample() next_state, reward, terminated, truncated, _ = env.step(action.cpu().numpy()) done = terminated or truncated log_prob = dist.log_prob(action).unsqueeze(0) entropy += dist.entropy().mean() log_probs.append(log_prob) values.append(value) rewards.append(torch.tensor([reward], dtype=torch.float, device=device)) masks.append(torch.tensor([1-done], dtype=torch.float, device=device)) state = next_state if done: print('Iteration: {}, Score: {}'.format(iter, i)) break next_state = torch.FloatTensor(next_state).to(device) next_value = critic(next_state) returns = discounted_rewards(next_value, rewards, masks) log_probs = torch.cat(log_probs) returns = torch.cat(returns).detach() values = torch.cat(values) advantage = returns - values actor_loss = -(log_probs * advantage.detach()).mean() critic_loss = advantage.pow(2).mean() optimizerA.zero_grad() optimizerC.zero_grad() actor_loss.backward() critic_loss.backward() optimizerA.step() optimizerC.step()