---
title: 'Homework 7: Reinforcement Learning'

---

# Homework 7: Reinforcement Learning
:::info
**Due December 4th, 2025 at 10 PM on Gradescope!**
:::

# Premise: Important Must Read ;)


In Bruno's **final challenge**, he needs to teach his underground supercomputer how to handle a variety of scenarios on its own. Your job is to write a handful of algorithms to reinforce good learning behavior into BSC (Bruno Super-Computer). From space ship landing to driving a mountain car, Bruno needs to ready to handle any scenario before he takes on CPax head on! Help a bear out :)

## Assignment Overview:
In this assignment we will be using OpenAI's gym to try out some techniques to train agents in the Cartpole Environment.

## Assignment Goals
In this assignment you are going to implement the Deep-Q, REINFORCE and REINFORCE with BASELINE training algorithms. You'll also get familiar with the gym library through the Cartpole environment.

## Getting Started
Please click [here](https://classroom.github.com/a/AcZesriD) to get the stencil code.

:::warning
**Here is the documentation for the [gymanisum docs](https://gymnasium.farama.org/)**
:::
:::danger
**Do not change the stencil except where specified.** You are welcome to write your own helper functions, however, changing the stencil's method signatures **will** break the autograder
:::

## Roadmap

### 1. Deep-Q Learning

#### Motivation
Q-Learning is the method of training an agent to learn so-called "q values" associated with taking an action when in a particular state. A q-value is a **proxy for reward** that also incorporates predicted future returns as well. To maximize cumulative reward, it makes sense to take the action with the highest predicted q-value at every timestep, which becomes our **policy**. Refer to the lecture slides for more details and a formal statement of the problem.

### Implementation

Head over to `deep_q.py`, here you will design your net for Deep Q learning

:::info
__Task 1.1 [DeepQModel Forward Pass]:__ Fill out the `__init__` and `call` methods for DeepQModel. The model should take a state/batch of states as input and output the predicted q-value of taking each action.

We recommend constructing an MLP using `tf.keras.Sequential`.

You should initialize your _target model_ to be a copy of your prediction model in the constuctor.

You'll need to build your prediction model $Q$ (check out the `.build` hint below), make a copy of $Q$ to be your target model, then copy the weights to the new target model.

You may find `tf.keras.models.clone_model`, `.set_weights`, and `.get_weights` helpful for this part!
:::

:::warning
For this assignment, optimizers should be created in the model `__init__`, you can then access the optimizer when training by using `model.optimizer`.
:::

:::warning
**About `.build`**

Tensorflow doesn't actually make the weights for a layer until it is called on an input. This way, it can stay generalizable to several different inputs (think about how dense layers don't need to have an input shape specified!)

We can force TF to build weights by calling `<insert model/layer_name>.build(input_shape=<insert shape iterable>)`
:::

:::success
__Note:__ In our case, the forward pass should take in a state and return the predicted reward __for each action__. The Q-value for a particular action can then be acquired by _indexing the output vector_.
:::

:::warning
An equally valid way to implement the network is to _append the action to the state as input_ and have the network return the predicted q-value of the state action pair. You __should not__ use this implementation for this assignment since it will not play nice with the autograder but feel free to come back to this assignment later on and make the necessary adjustments to train the network in this way.
:::
:::info
__Task 1.2 [DeepQModel Loss]:__ Implement the loss function for a DQN.
- Keep in mind each batch is a list of 5 items: _batch_states, batch_actions, batch_rewards, batch_next_states, batch_done_. Each is a list of length batch_size. Each example in the batch is independent so each example is non-consecutive and can be from different episodes entirely.
- The loss function _per example_ is as follows:
$$\mathcal{L}(Q, \hat{Q}, s, a, r, s', d) = \left(Q(s,a)-\left[r + (1-\mathbb{1}_{d})\cdot \gamma \cdot\max_{a'}\hat{Q}(s',a')\right] \right)^2$$
Where $s$ is the current state, $a$ is the action taken, $s'$ is the next state, $r$ is the reward, $\gamma$ is the discount factor (usually .99), $d$ is a boolean which is True if the game terminated after this action, false otherwise. $Q$ is the prediction network and $\hat{Q}$ is the target network.
- $\mathbb{1}_{d}$ is 1 if d is True and 0 otherwise. If the state is terminal, you only need to consider the reward at that state and not any future returns.

- As typical, you should vectorize this function and __return the mean loss over each example in the batch.__
:::
:::success
__Verification:__
You can check your implementation by running the following command:

python local_tests.py DEEP_Q

:::

Now that we have a model, we need a way to train it!

:::success
__Note:__ `assignment.py` has several functions, many of which you will implement, some of which you won't have to worry about. Here is a breakdown of the different functions:
- `visualize_data` and `visualize_episode` are utility methods we will use to visualize your trained agents and rewards acquired over the course of training. You don't have to worry about these, but you should look at their function signatures if you want to use them during/after training.
- `discount` and `train_trajectory` are helper functions for REINFORCE and REINFORCE with BASELINE. You'll implement these later and use them to train with those algorithms but for now you can skip past them.
- `train_reinforce_episode` and `train_deep_q_episode` should train the passed model on the passed environment for 1 episode. `train_deep_q_episode` has extra args for `batch_size`, `memory` and `epsilon` which are applied in the training algorithm for deep Q learning. You will be filling out `train_deep_q_episode` shortly.
- `train` is used to run a single training episode, it will call either `train_reinforce/deep_q_episode` depending on what model was passed in.
:::
:::info
__Task 1.3 [Prepare to train a DeepQModel part 1]:__ Implement `train_deep_q_episode`.

`train_deep_q_episode` should
1. Reset the environment and prepare any variables you need
2. Simulate a whole episode, append the information from each step into the memory bank. Each item in `memory` should be a `(state, action, reward, next_state, done)` tuple.
3. For n batches (around 10) train the model by:
- Randomly choosing `batch_size` many examples from the memory
- Convert the batch into to tensors and arrays
- Since each item in memory has a different datatype, you'll need to split the values of your batch before converting to tensors or arrays.
- Compute the loss and update the weights of the model
4. Update the target network parameters
5. Return the sum total of rewards earned during the episode simulation and the updated memory bank
:::

:::success
__Note:__ When simulating an episode, you can interact with the environment with
- `env.step(action)` and `env.sample()`
:::
:::info
__Task 1.4 [Prepare to train a DeepQModel part 2]:__ Implement train.

`train` should
1. Check which model is being trained (use `isinstance` for this)
2. If training a `DeepQModel` and the memory is empty, initialize a memory of around 50 samples using random actions. Before calling `train_deep_q_episode`, be sure to **truncate the memory to 1000 examples**.
3. If training either `Reinforce` or `ReinforceWithBaseline`, just call `train_reinforce_episode`
4. Return both the total reward and memory! This can be `None` for `Reinforce` and `ReinforceWithBaseline`
:::

:::success
**You should know:** Truncating the memory might at first seem like a bad idea since we are losing out on training data. This method has lots of merits however. Consider that our oldest examples are likely to contain worse state-action pairs that are not helpful to the training process. We'd prefer not to keep learning from bad examples and focus on state-action pairs that are closer to optimal.
:::
:::info
__Task 1.4 [Train a DeepQModel]:__ You can run `main` by using the command line `python assignment.py --model DEEP_Q`. You can adjust the main function to tweek the `num_episodes`, how often episodes are visualized, etc. but the default values should work well. Ensure you can train a `DeepQModel` with average reward > 110.
:::
---
### 2. REINFORCE
#### Motivation
REINFORCE is a policy gradient algorithm so instead of trying to learn the reward values it directly optimizes for the policy function.
:::info
__Task 2.1 [Reinforce initialization and call]:__ Implement the init and `call` method for `Reinforce`, we recommend a small MLP for this.
:::

:::info
__Task 2.2 [Reinforce loss function]:__ Implement `loss_func` for `REINFORCE`.

The loss is given by
$$\mathcal{L}(s_t, a_t) = -\sum_{t=0}^T\log(\pi_\theta(a|s))* G_t $$
Where $\pi_\theta(a|s)$ is the predicted log probability of taking action, $a$, from state, $s$ by the policy network $\pi_\theta$ with parameters $\theta$. $G_t$ is the discounted rewards for the trajectory starting at timestep $t$.

:::
:::success
__Verification:__
You can check your implementation by running the following command:

python local_tests.py REINFORCE

:::
:::success
__Note:__ A hint when implementing `Reinforce.loss_func`

- You'll need to call the model and use `gather_nd` to get the probabilities associated with each action
:::

:::info
__Task 2.3 [Prepare to train REINFORCE, discount]:__ Head over to `assignment.py` here you will find the `discount` function. Given a list of rewards, $r$, and discount rate $\gamma$, you should return a list of length $t$ of the discounted rewards from each timestep $G_t$, where:

$$G_t = \sum_{i=0}^{T-t} \gamma^i*r_{i+t}$$

An example of expected behavior is provided in the handout.
:::

:::info
__Task 2.4 [Prepare to train REINFORCE, generate_trajectory]:__

Here you will simulate 1 episode using the model to generate the policy, you should return lists of all states, actions and rewards experienced during the episode.
:::

:::warning
__Note:__ Recall that policy gradient methods output a _probability distribution_ over the action space and that this distribution __is a representation of our policy__. Thus, you should sample from this distribution rather than take a maximum predicted action.
:::

:::info
__Task 2.5 [Train REINFORCE]:__ Put your model, `generate_trajectory` and `discount` together to implement `train_reinforce_model` which should train on one simulated episode.

Be sure to adjust `train` to call `train_reinforce_model` when using `Reinforce` and `ReinforceWithBaseline` models if you didn't do so in Task 1.4.
:::

:::success
You should now be able to train REINFORCE by running `python assignment.py --model REINFORCE`
:::

### 3. REINFORCE with BASELINE

:::info
__Task 3.1 [ReinforceWithBaseline]:__ At this point, all you have to do are fill in the methods for `ReinforceWithBaseline` then you can run `python assignment.py --model REINFORCE_BASELINE` to train using the new model. We leave the details for this, the last task for 1470, to you (with a few helpful hints below).
:::
:::success
__Verification:__
You can check your implementation by running the following command:

python local_tests.py REINFORCE_BASELINE

:::

:::success
__Hints:__
- In our loss functions, instead of directly using discounted returns, we use the so-called advantage function given by $$A^\pi(s_t, a_t) = G_t - V_{\pi_\theta}(s_t)$$ where $V_{\pi_\theta}$ is the value network. There are many possible ways to compute advantage that give rise to other algorithms. For instance, if you use the expected returns instead of the returns for a single trajectory, you arrive at Advantage Actor Critic (A2C):
$$A^\pi(s_t, a_t) = r + \gamma V_{\pi_\theta}(s_{t+1}) - V_{\pi_\theta}(s_t)$$
- The Actor loss for a single step is given by $$-(A^\pi(s_t, a_t)\cdot\log(\pi_\theta(a|s_t))$$ where $\pi_\theta$ is the actor network with parameters $\theta$
- Be sure to use `tf.stop_gradient` around the $A^{\pi_\theta}$ when computing actor loss, if you do not, then the actor loss gradient will continue to propagate through the critic network. Consult the box below if you are unsure why this is necessary
- The critic loss for a single step is given by $$[A^\pi(s_t, a_t)]^2$$
- As with all loss functions in this course, you should return one value, the sum of the Critic and Actor losses. In pseudo-code, we usually write these two loss functions (and updates) as different steps. In practice, you can combine them into a single loss function so long as you compute each individual loss function in an appropriate gradient tape environment.
:::

:::success
__You should know:__
When you completed BERAS, the gradient method may have seem overly generalized and it may have been hard to imagine cases in which a search algorithm would have been useful. The above loss function is a perfect example of why the generalization is important. With a single GradientTape we can compute every gradient we need across both networks at once.

Think back to your gradient method and consider how it would have handled the $A^{\pi}$ term in the Actor loss. Since $A^\pi$ is parameterized, its weight gradients would have been computed and its weights then updated had we not used `tf.stop_gradient`.
:::

## Running the Assignment

The assignment uses command-line arguments to control training, testing, and submission. Below are the available parameters and how to use them. You can run the assignment using

python assignment.py [-help] 
--model {REINFORCE,REINFORCE_BASELINE,DEEP_Q} 
[--env ENV] 
[--watch] 
[--load-path LOAD_PATH]
[--save-path SAVE_PATH] 
[--lr LR] 
[--num-episodes NUM_EPISODES] 
[--submit]

### Command Line Parameters

- `--model`: Model type to train. **Required.** Choices: `REINFORCE`, `REINFORCE_BASELINE`, `DEEP_Q`
- `--num-episodes`: (Optional) the number of episodes to train the agent on the scenario for. Default: `650`
- `--env`: (Optional) Gym environment name. Default: `CartPole-v1`
- If you do not specify one, then it will prompt you to seelct one
- `--lr`: (Optional) Learning rate for training. Default: `0.001`
- `--watch`: (Optional) Enable watch mode to render a trained agent **after** training
- `--save-path`: (Optional) Path to save model weights after training (optional - best weights are saved automatically)
- `--load-path`: (Optional) Path to load pre-trained model weights (optional - best weights are loaded automatically in watch mode)
- `--submit`: (Optional) Save training plots to the `submission/` folder for grading

### Usage Examples

**Basic training:**

python assignment.py --model REINFORCE --num-episodes 650

This will then prompt you select an environment if you do not specify one as

==================================================
Available Environments:
1. CartPole-v1
2. LunarLander-v3
3. MountainCar-v0
4. Acrobot-v1
==================================================

Select an environment (enter number):

**Train on a specific environment:**

python assignment.py --model DEEP_Q --env LunarLander-v2 --num-episodes 1650

**Train with custom learning rate:**

python assignment.py --model REINFORCE_BASELINE --num-episodes 650 --lr 0.002

**Watch your best trained model (simplified - no path needed):**

# First train the model
python assignment.py --model DEEP_Q --env CartPole-v1

# Then watch it perform (automatically loads best weights)
python assignment.py --model DEEP_Q --env CartPole-v1 --watch

**Save model weights to a custom location:**

python assignment.py --model DEEP_Q --save-path ./my_custom_weights/dqn_model

**Load and watch a model from a custom path:**

python assignment.py --model DEEP_Q --load-path ./my_custom_weights/dqn_model --watch

## Submission

To submit your assignment, you need to generate training plots for each model you've implemented.

### Generating Submission Files

Run the following commands to train each model and save the training plots (you can specify more commands if you are running a specific env, lr, or num-episodes):

python assignment.py --model REINFORCE --submit
python assignment.py --model REINFORCE_BASELINE --submit
python assignment.py --model DEEP_Q --submit

Each command will:
1. Train the model for 650 episodes
2. Display training progress with a progress bar
3. Save a plot to `submission/<MODEL>_train_plot.png`
4. Display the final training plot

The `submission/` folder will contain:
- `REINFORCE_train_plot.png`
- `REINFORCE_BASELINE_train_plot.png`
- `DEEP_Q_train_plot.png`

Submit the entire `submission/` folder along with your code.

### What the Plots Show

Each training plot contains two subplots:
- Episode rewards over time with a 50-episode moving average
- Training loss over time

These plots help demonstrate that your implementation is learning correctly. You will need to **submit these for each model type**!

### Submission Instructions

We recommend submitting your code via GitHub.

**Minimum Required Files:**
To ensure the autograder runs correctly, your submission **must** include the following files and structure:
- `__init__.py` (at the root of your repository)
- `src/` directory (containing `train.py`, `deep_q.py`, etc.)
- `tests/` directory

**Important:** The top-level `__init__.py` file is required for the autograder to recognize your code as a package. If it is missing, the autograder will fail with an import error.

If you choose to submit manually by uploading a zip file, ensure that your zip file includes the top-level `__init__.py` and maintains the directory structure listed above.

## Additional Features

Cartpole is a very easy environment so you will likely need to run more episodes when training on other environments. That said, it can be a lot of fun to watch a network learn a more complicated environment with better graphics!

### Local Testing
We have provided a set of minimal tests to help you check your implementation as you go. These tests check for basic functionality like correct output shapes and types, but do not guarantee that your model will converge or solve the environment.

To run the local tests:

python local_tests.py

These tests cover initialization, `call` output shape, and `loss_func` execution, but not correctness. That is for gradescope lol :)

If you encounter `NotImplementedError`, it means you haven't implemented that part of the assignment yet.

You can run specific tests using:

python local_tests.py DEEP_Q
python local_tests.py REINFORCE
python local_tests.py REINFORCE_BASELINE

## Thank you

Congratulations you have completed your last assignment for CSCI 1470!