# Diffusion-QL: Reproduction & Innovation (CQL Integration)

**Author:** Shaimon Rahman  
**Project Type:** Research Reproduction & Extension  

This repository contains the source code and experimental results for the project:  
**“Reproduction and Extension of Diffusion Policies for Offline Reinforcement Learning”**

It builds upon the official implementation of [Wang et al. (2023)](https://github.com/Zhendong-Wang/Diffusion-Policies-for-Offline-RL) by integrating **Conservative Q-Learning (CQL)** regularization to address overestimation bias in diffusion-based policies.


---

## Key Features

1.  **Modified Source Code:** Full implementation of Diffusion-QL with added support for conservative regularization.
2.  **Innovation (CQL Penalty):** A custom implementation of the CQL penalty term added to the critic loss function to handle out-of-distribution (OOD) actions.
3.  **Reproducible Notebook:** A complete `replicate_run.ipynb` that runs the baseline reproduction, ablation studies, and the new innovation experiments.

---

##  Methodology & Innovation

This project extends the original Diffusion-QL framework by integrating **Conservative Q-Learning (CQL)** (Kumar et al., 2020). The standard Diffusion-QL method can sometimes suffer from Q-value overestimation when the diffusion model generates actions far from the training distribution.

To mitigate this, I introduced a conservative penalty term into the critic loss:

### Implementation Details

#### 1. `agents/ql_diffusion.py`
The critic loss was modified to include the CQL regularizer:
```python
# [Innovation: Conservative Q-Learning Penalty]
# Calculate Conservative Penalty
cql_term = self.cql_alpha * (torch.logsumexp(q_cand, dim=1).mean() - q_data)
critic_loss = critic_loss + cql_term
```

#### 2. `main.py`
* Added CLI arguments `--cql_alpha` and `--cql_n_action` to control the penalty strength dynamically.
* Added logging for `"CQL Term"` to track the penalty magnitude during training.

---

## How to Run

### Prerequisites
Install the required dependencies:
```bash
pip install -r requirements.txt
*Note: Requires PyTorch, MuJoCo, and D4RL.*

### 1. Running the Innovation (CQL Experiments)

To reproduce the CQL $\alpha$-sweep experiments:

```bash
# Alpha = 0.0 (Baseline)
python run_with_d4rl.py --env_name hopper-medium-replay-v2 --algo ql \
  --ms offline --T 5 --batch_size 128 --num_steps_per_epoch 200 \
  --lr_decay --early_stop --cql_n_action 10 --cql_alpha 0.0 --exp hopper-mr-ql-cql0p0-dev

# Alpha = 0.2 (Conservative)
python run_with_d4rl.py ... --cql_alpha 0.2 --exp hopper-mr-ql-cql0p2-dev

# Alpha = 0.5 (Highly Conservative)
python run_with_d4rl.py ... --cql_alpha 0.5 --exp hopper-mr-ql-cql0p5-dev
```
### 2. Using the Notebook

The file `replicate_run.ipynb` is a self-contained lab that includes:

* Baseline reproduction on `hopper-medium-replay-v2`.
* Ablation studies on diffusion steps ($T \in \{2, 5, 10\}$).
* Comparisons between Behavior Cloning (BC) and Diffusion-QL.
* Visualization of results and training curves.

---

## Original Paper & Abstract

**Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning** Zhendong Wang, Jonathan J. Hunt, and Mingyuan Zhou (ICLR 2023)  
[arXiv:2208.06193](https://arxiv.org/abs/2208.06193)

> **Abstract:** Offline reinforcement learning (RL) aims to learn an optimal policy using a previously collected static dataset. Standard RL methods often perform poorly due to function approximation errors on out-of-distribution actions. [...] In this paper, we propose representing the policy as a diffusion model, a recent class of highly-expressive deep generative models. We introduce Diffusion Q-learning (Diffusion-QL) that utilizes a conditional diffusion model for behavior cloning and policy regularization.

### Citation

If you use the original Diffusion-QL code, please cite:

```bibtex
@article{wang2022diffusion,
  title={Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning},
  author={Wang, Zhendong and Hunt, Jonathan J and Zhou, Mingyuan},
  journal={arXiv preprint arXiv:2208.06193},
  year={2022}
}