AI systems research / Course collaboration / 2026

Can reinforcement learning
improve molecular docking?

I built an end-to-end GRPO-style post-training loop for DiffDock. One small split improved. A second did not. The value of the project is the working system—and the evidence showing where it breaks.

My focusDiffDock experiment implementation

Built withPython · PyTorch · DiffDock · GRPO

Evidence statusPreliminary local experiments

Inspect the repository

01 / The research question

A narrower testbed.
A difficult transfer.

The broader course collaboration explored reinforcement learning for molecular generation. My case study focuses on a narrower question: could the same group-relative idea post-train a pretrained diffusion model for fixed-ligand docking?

DiffDock predicts a ligand pose through translation, rotation, and torsion diffusion. That made it a practical intermediate system for testing reward-based updates without claiming to solve full molecular generation.

Scope note: the broader course project was collaborative. This page describes the DiffDock extension represented in my repository and does not claim sole authorship of teammates’ work.

02 / System

How the system works

Each update connects stochastic docking, a structural reward, a group-relative comparison, and DiffDock’s own training losses. The checkpoint then returns to the original inference path so baseline and updated models are judged the same way.

Generate

DiffDock samples four poses for a protein–ligand complex.

Reward

RMSD scores each pose against the known ligand position.

Compare

Rewards become normalized advantages within the group.

Update

A native DiffDock loss supplies a differentiable surrogate signal.

Export

The updated score model returns to DiffDock checkpoint format.

Evaluate

The same pose-generation and RMSD pipeline tests the checkpoint.

The update signal

sθ = −(Ltranslation + Lrotation + Ltorsion)

Group size 4 · score batch size 2 · learning rate 1×10⁻⁷

03 / Role

My contribution

I implemented the DiffDock model loading and scoring path, native-loss batch construction, the GRPO-style update, batched gradient accumulation, checkpoint export, and the evaluation workflow. I also added tests and memory-oriented batching fixes as the experiment moved from a dry run to real complexes.

Training loop

Grouped on-policy rollouts, RMSD rewards, normalized advantages, and a clipped surrogate objective.

Systems work

Score batching and callable batch generation kept the native DiffDock loss path within available memory.

Evaluation

Exported checkpoints were loaded back into DiffDock and compared through the same pose and RMSD pipeline.

Contribution boundaries are based on the repository’s implementation history. The numerical findings below come from a local experiment report draft and should be read as preliminary, not as a published benchmark.

04 / Findings

Two seeds. Two different stories.

The first split suggested a meaningful improvement. The cleaner second split did not reproduce it. Showing both is more useful than presenting the best run alone.

Seed 44

Promising signal

Improved

BaselineAfter one update

Coverage0.40 → 0.60

Strict Success@10.40 → 0.60

Mean RMSD7.49 Å → 3.03 Å

Severe RMSD outliers also fell from 0.250 to 0.083.

Seed 45

Failed to generalize

Unstable

BaselineAfter one update

Coverage1.00 → 0.60

Strict Success@10.80 → 0.40

Mean RMSD9.17 Å → 8.94 Å

Coverage fell, strict Success@1 halved, and the RMSD change was marginal.

Coverage and Success@1 are proportions; higher is better. RMSD is measured in ångströms; lower is better. These small-split results do not establish robust model improvement.

05 / Friction

What held the work back

Update stability. Repeated or poorly regularized updates could reduce generation coverage or move poses outside the binding pocket.

Small evaluation slices. The result is too sensitive to seed and complex selection for broad performance claims.

Surrogate mismatch. Optimizing DiffDock’s native loss is tractable, but it is not the same as differentiating through the full sampling trajectory.

Compute constraints. Score batches of two and gradient accumulation were engineering necessities, not an idealized experimental setup.

06 / Next action

What I would test next

StabilizeCompare one-step updates, stronger regularization, and gradient clipping before scaling.
FilterIdentify unstable complexes and distinguish data failures from optimization failures.
RepeatRun more seeds with a frozen evaluation protocol and report uncertainty, not only point estimates.
Scale lastOnly expand the benchmark after the improvement survives those checks.

The current conclusion is intentionally narrow: the end-to-end post-training system works, but the improvement is not yet reliable.

Return to selected work

A narrower testbed.A difficult transfer.