The key new idea is:

DQN code = trains the model
PyTorch Profiler = measures what takes time during training

1. Import Gymnasium

import gymnasium as gym

Gymnasium provides the CartPole environment.

env = gym.make("CartPole-v1")

creates the game. The agent sees a state containing 4 values:

[cart position,
 cart velocity,
 pole angle,
 pole angular velocity]

It can choose 2 actions:

0 → Left
1 → Right

2. Import PyTorch

import torch

PyTorch is used for:

Tensors
Neural networks
Calculating loss
Backpropagation
Updating model parameters
Profiling

3. Import neural network tools

import torch.nn as nn

nn provides neural-network layers. For example:

nn.Linear()
nn.ReLU()
nn.Sequential()

4. Import optimizer

import torch.optim as optim

The optimizer updates the model’s weights. You’re using:

optim.Adam()

5. Create CartPole

env = gym.make("CartPole-v1")

This creates the environment. Think of it as:

       Pole
        │
        │
   ─────┴─────
      Cart

← Move      Move →

The goal is to keep the pole balanced for as long as possible.

6. Create the DQN model

model = nn.Sequential(
    nn.Linear(4, 32),
    nn.ReLU(),
    nn.Linear(32, 2)
)

This is the neural network. Its structure is:

Input
  ↓
4 values
  ↓
Linear(4 → 32)
  ↓
ReLU
  ↓
Linear(32 → 2)
  ↓
2 Q-values

For example:

State
[0.1, -0.2, 0.03, 0.15]

       ↓

     DQN

       ↓

[4.2, 6.7]

Meaning:

Q(left)  = 4.2
Q(right) = 6.7

The agent would choose right.

7. Create optimizer

optimizer = optim.Adam(
    model.parameters(),
    lr=0.001
)

`model.parameters()`

Gets all trainable weights from the network.

`lr=0.001`

Learning rate. It controls how much the weights change after each training step.

8. Start the profiler

This is the major new part.

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU
    ],
    record_shapes=True
) as prof:

This tells PyTorch:

“While the following code is running, monitor and record what the program is doing.”

What is a profiler?

A profiler helps answer questions like:

What is taking the most CPU time?

Which PyTorch operation runs most often?

Which operation is expensive?

What tensor shapes are being processed?

Think of it like a performance monitoring tool for your ML code.

9. CPU profiling

torch.profiler.ProfilerActivity.CPU

This tells the profiler to monitor CPU operations. You’re currently not profiling GPU/CUDA operations. So:

Your program
     │
     ▼
PyTorch
     │
     ▼
CPU operations
     │
     ▼
Profiler records them

10. `record_shapes=True`

record_shapes=True

This tells the profiler to record the shapes of tensors involved in operations. For example, it can help you see something like:

Linear operation
Input shape: [4]
Output shape: [32]

This is useful when debugging performance.

11. Start training

Everything inside:

with torch.profiler.profile(...) as prof:

gets profiled. Your training starts here:

for episode in range(5):

You’re running only 5 episodes. This is intentionally small because profiling adds overhead.

12. Reset environment

state, info = env.reset()

Starts a new CartPole episode. You get the initial state. Example:

state =
[0.01, -0.02, 0.03, 0.04]

13. Maximum 100 steps

for step in range(100):

The agent can take up to 100 actions in each episode. So potentially:

5 episodes × 100 steps

up to around 500 environment steps, although an episode can terminate earlier.

14. Convert state to tensor

state_tensor = torch.tensor(
    state,
    dtype=torch.float32
)

Gymnasium gives you numerical values. PyTorch’s neural network expects tensors. So:

Gym state
    ↓
[0.01, -0.02, 0.03, 0.04]
    ↓
PyTorch Tensor

15. Forward pass

q_values = model(state_tensor)

This is called the forward pass. The state goes through the neural network:

State
  ↓
Linear
  ↓
ReLU
  ↓
Linear
  ↓
Q-values

Example:

[0.01, -0.02, 0.03, 0.04]
                ↓
          Neural Network
                ↓
           [2.4, 3.7]

16. Select action

action = q_values.argmax().item()

argmax() finds the largest Q-value. If:

q_values = [2.4, 3.7]

then:

3.7 → index 1

So:

action = 1

means:

Move right

17. Take the action

next_state, reward, terminated, truncated, info = env.step(
    action
)

The agent sends the action to CartPole. The environment returns:

next_state
reward
terminated
truncated
info

For example:

Action → Right

       ↓

CartPole

       ↓

New state
Reward = 1

18. Calculate target

target = torch.tensor(
    reward,
    dtype=torch.float32
)

Here you’re using the immediate reward as the target. If:

reward = 1

then:

target = 1

Important

This is not a complete DQN target calculation. A proper DQN would normally use something like:

target = r + \gamma \max Q(s',a')

when the episode hasn’t ended. Your earlier code actually included the future Q-value. This simplified version does not.

19. Calculate loss

loss = (q_values[action] - target) ** 2

Suppose:

q_values[action] = 2.5
target = 1

Then:

loss = (2.5 - 1)²
     = 2.25

The model is trying to make:

Predicted Q-value
        ↓
closer to
        ↓
Target

20. Clear old gradients

optimizer.zero_grad()

PyTorch accumulates gradients by default. So before calculating new gradients, we clear the previous ones.

Old gradients → remove
New gradients → calculate

21. Backward pass

loss.backward()

This performs backpropagation. It calculates gradients showing how the model’s parameters contributed to the loss. Conceptually:

Loss
 ↓
Backward propagation
 ↓
Gradients
 ↓
Which weights should change?

22. Update model

optimizer.step()

Adam uses the calculated gradients to update the neural network’s parameters. So:

Forward pass
     ↓
Q-value
     ↓
Loss
     ↓
Backward pass
     ↓
Gradients
     ↓
Adam
     ↓
Updated weights

That’s the actual learning process.

23. Update state

state = next_state

The next state becomes the current state.

State A
   ↓
Action
   ↓
State B
   ↓
State B becomes current state

Then the process repeats.

24. Stop if episode ends

if terminated or truncated:
    break

If the pole falls or the environment reaches its limit, stop that episode. Then the next episode starts.

25. Print profiler results

After training finishes:

print(
    prof.key_averages().table(
        sort_by="cpu_time_total",
        row_limit=10
    )
)

This is where you see the profiling information.

`prof.key_averages()`

Groups similar operations together. Instead of showing every single call separately, it summarizes operations. For example:

aten::linear
aten::relu
aten::add
aten::mul
aten::pow

`.table()`

Formats the results as a table.

`sort_by="cpu_time_total"`

Sorts operations according to their total CPU execution time. So the most expensive operations appear near the top.

`row_limit=10`

Only show the top 10 operations. You may see output conceptually similar to:

---------------------------------------------------------------
Name              CPU total    CPU time avg    Calls
---------------------------------------------------------------
aten::linear      ...          ...              ...
aten::addmm       ...          ...              ...
aten::relu        ...          ...              ...
aten::pow         ...          ...              ...
aten::sub         ...          ...              ...
---------------------------------------------------------------

The exact numbers depend on your computer and run.

26. Close environment

env.close()

This releases resources used by the CartPole environment.

The Important Difference From Your Previous Code

Your previous code was mainly about:

How does DQN learn?

This code adds:

How expensive is each operation?

So you now have two concepts:

DQN

State
 ↓
Neural Network
 ↓
Q-values
 ↓
Action
 ↓
Reward
 ↓
Loss
 ↓
Backpropagation
 ↓
Update model

Profiler

Training code
      ↓
PyTorch Profiler
      ↓
Record operations
      ↓
Measure CPU time
      ↓
Print performance table

​1. Import Gymnasium

​2. Import PyTorch

​3. Import neural network tools

​4. Import optimizer

​5. Create CartPole

​6. Create the DQN model

​7. Create optimizer

​model.parameters()

​lr=0.001

​8. Start the profiler

​What is a profiler?

​9. CPU profiling

​10. record_shapes=True

​11. Start training

​12. Reset environment

​13. Maximum 100 steps

​14. Convert state to tensor

​15. Forward pass

​16. Select action

​17. Take the action

​18. Calculate target

​Important

​19. Calculate loss

​20. Clear old gradients

​21. Backward pass

​22. Update model

​23. Update state

​24. Stop if episode ends

​25. Print profiler results

​prof.key_averages()

​.table()

​sort_by="cpu_time_total"

​row_limit=10

​26. Close environment

​The Important Difference From Your Previous Code

​DQN

​Profiler