documenting DPO scoring for infeasible prompt sequences notes, editing readme

Author: Drew Prinster {External}

README.md

@@ -1,6 +1,6 @@
 # Conformal Policy Control
 
-Code for ["Conformal Policy Control"](https://arxiv.org/abs/2603.02196) (ICML 2026 spotlight paper): a framework for enabling AI agents to automatically determine their own "zone of competence," where we can guarantee their behavior will respect a user's risk tolerance, $\alpha$.
+Code for ["Conformal Policy Control"](https://arxiv.org/abs/2603.02196) (ICML 2026 spotlight paper): a framework for enabling AI agents to automatically determine their own "zone of competence," where we can place guarantees on their behavior respecting a user's risk tolerance, $\alpha$.
 
 By Drew Prinster, Clara Fannjiang, Ji Won Park, Kyunghyun Cho, Anqi Liu, Suchi Saria, and Samuel Stanton.
 
@@ -18,26 +18,28 @@ If you use this code, please our paper:
 
 ## Overview
 
-This project develops **Conformal Policy Control (CPC)**: a method for iteratively improving a language model policy while maintaining formal guarantees on the risk (e.g., rate of infeasible or unsafe outputs) over time. The key idea is to constrain each optimized policy's likelihood ratios relative to a safe reference policy, with the constraint level calibrated via conformal prediction so that risk stays below a user-specified level alpha.
+This project develops **Conformal Policy Control (CPC)**: a method for iteratively improving a language model policy while maintaining formal guarantees on the risk (e.g., rate of infeasible or unsafe outputs) over time. The key idea is to constrain each optimized policy's likelihood ratios relative to a safe reference policy, with the constraint level calibrated via CPC so that risk stays below a user-specified level, $\alpha$.
 
 ![CPC animation: search and sampling](visuals/Animation_alpha0.5_betaHat10.400000_CPCsearchTrue_samplingTrue_proposalsTrue.gif)
 
 The repository contains four sets of experiments:
 
-- **`cpc_llm/`** : The main CPC pipeline for LLMs, applied to the Ehrlich function protein discovery task ([Chen, et al. 2025](https://arxiv.org/abs/2410.22296)). Pre-trains a LM on data from a genetic algorithm, then iteratively generates and scores new samples, trains optimized policies (SFT, DPO, or MARGE), and uses CPC to ensure the improved policies satisfy safety constraints.
-- **`cbo/`** : Constrained Bayesian optimization experiments (in paper appendix). Compares CPC to classic conservative optimization. Simplest initial entrypoint to CPC, runs on single CPU.
+- **`cpc_llm/`** : **The main CPC pipeline for LLMs**, applied to the Ehrlich function synthetic protein discovery task ([Chen, et al. 2025](https://arxiv.org/abs/2410.22296)). Pre-trains a language model on data from a genetic algorithm, then iteratively generates and scores new samples, trains optimized policies (SFT, DPO, or MARGE), and uses CPC to ensure the improved policies satisfy safety constraints.
+- **`cbo/`** : Constrained Bayesian optimization experiments (in paper appendix). Compares CPC to classic conservative optimization. **This is a more accessible initial entrypoint to CPC code (runs on a single CPU).**
 - **`constrained_AL/`** : CPC constrained active learning with Gaussian process surrogates, applied to tabular regression benchmarks. 
 - **`QA_expts/`** : Generalized conformal risk control (gCRC) for LLM factuality, controlling false discovery rate (a non-monotonic loss) on medical QA dataset of GPT-3.5-Turbo responses.
 
 ## Setup
 
 Requires Python >= 3.10 and [uv](https://docs.astral.sh/uv/).
 
+To install all dependencies (including the `cpc-llm` package in editable mode) and dev tools (pytest), and then activate the environment, run
+
 ```bash
 uv sync --group dev
-```
 
-This installs all dependencies (including the `cpc-llm` package in editable mode) and dev tools (pytest).
+source .venv/bin/activate
+```
 
 ## Running the CPC-LLM pipeline
 
@@ -52,6 +54,9 @@ cpc-llm --config-name=cpc_llm \
   conformal_policy_control.alpha=0.6 initial_seed=0 last_seed=0 \
   local_output_dir=/path/to/local \
   parent_output_dir=s3://bucket/path
+
+# With slurm (and changing resource params in run_cpc_llm.sh), run many parallel GPU jobs to reproduce paper's Fig 6 via
+bash submit_cpc_llm_expts.sh "0.4,0.6,0.8,1.0" 0 29
 ```
 
 
@@ -60,15 +65,16 @@ cpc-llm --config-name=cpc_llm \
 
 | Parameter | Description |
 |-----------|-------------|
-| `conformal_policy_control.alpha` | Risk level (e.g., 0.1 for 10% constraint violation rate) |
+| `conformal_policy_control.alpha` | Risk level (e.g., 0.4 for 40% constraint violation rate) |
 | `num_sft_rounds` / `num_dpo_rounds` / `num_marge_rounds` | Number of training iterations per method |
+| `initial_seed` / `last_seed` | Initial / last random seeds (inclusive) to run in loop for repeat trials |
 | `parent_output_dir` | S3 path for outputs (set to `null` for local-only) |
 | `local_output_dir` | Local path for outputs and model checkpoints |
 
 ### Important notes
 
 - **Storage**: The pipeline supports both local and S3 storage. Model checkpoints are copied to S3 and deleted locally after training. Set `parent_output_dir: "null"` to disable S3.
-- **SLURM**: Training, generation, and scoring jobs are launched as SLURM jobs. Configure via `slurm_args` sections in the config.
+- **SLURM**: Training, generation, and scoring jobs are launched as their own SLURM jobs. Configure via `slurm_args` sections in the config.
 - **Resuming**: The pipeline automatically resumes prior runs if launched with the same config. Use `--overwrite=True` to start fresh.
 - **GPU requirements**: Training uses DDP (single-node multi-GPU). You need ~4x the model size in GPU RAM for full-precision training.
 
@@ -94,6 +100,7 @@ cpc_llm/                  # Main CPC-LLM package (installable)
     infrastructure/       # File handling (local/S3), orchestration, SLURM
     train/                # SFT, DPO, MARGE training
     test_functions/       # Ehrlich benchmark utilities
+cbo/                      # Constrained Bayesian optimization experiments
 constrained_AL/           # Active learning experiments
 QA_expts/                 # Medical QA experiments
 notebooks/                # Visualization notebooks
@@ -103,3 +110,7 @@ tests/                    # Unit tests
 ## License
 
 See [LICENSE](LICENSE).
+
+## Known implementation notes
+
+- **DPO scoring for infeasible prompt sequences**: `DPOTrainerWithLogging` does not count a prompt→response transition from infeasible to feasible as a score improvement, which makes DPO training more permissive of infeasible outputs than originally intended. This does not invalidate the paper's experiments (its main effect is increasing the risk of the unconstrained policy), so it is kept as-is for reproducibility. See the full explanation in [`pref_tuning_trainer.py`](cpc_llm/src/cpc_llm/train/pref_tuning_trainer.py#L123-L141).

cpc_llm/src/cpc_llm/data/synthetic_dataset_formatter.py

@@ -425,6 +425,14 @@ def find_preference_pairs(cfg: DictConfig, df: pd.DataFrame) -> List[Dict[str, A
     # comparisons against NaN are always False, which would silently produce
     # zero pairs instead of treating the particle as worst-possible.
     scores_np = filtered_scores.numpy()
+    
+    ## Note on DPO scoring for infeasible prompt sequences: 
+    ## The following commented out line "scores_np = np.where(np.isnan(scores_np), np.inf, scores_np)" 
+    ## relates to the "Note on DPO scoring for infeasible prompt sequences" in ../train/pref_tuning_trainer.py.
+    ## Please see the note in that file.
+
+    # scores_np = np.where(np.isnan(scores_np), np.inf, scores_np)
+
     pynn_transformer = PyNNDescentTransformer(
         n_neighbors=cfg.n_neighbors, metric=cfg.distance_metric
     ).fit(filtered)

cpc_llm/src/cpc_llm/train/pref_tuning_trainer.py

@@ -119,6 +119,27 @@ def get_metrics_for_outputs(
                 num_repeated_input += 1
             particle = torch.FloatTensor(particle).unsqueeze(0)
             score = self.test_fn(particle).item()
+
+            ## Note on DPO scoring for infeasible prompt sequences:
+            ## The following if statement patches a NaN/None handling bug caused by inf scores
+            ## (for infeasible sequences) being recorded as NaN/None. The original intention was to increment 
+            ## num_decreased_score also in the case that the prompt was infeasible (i.e., input_scores[i]
+            ## is None or inf) and the response was feasible (i.e., score is not None nor inf). However, the 
+            ## current implemented behavior does *not* increment in this case, which essentially makes the
+            ## DPO training more "unsafe" / less incentivized to avoid infeasible outputs, which makes
+            ## for a setting where the effect of CPC "reining in" an unsafe model can be demonstrated. 
+            ## That is, this "bug" makes for a less standard DPO training, but does *not* invalidate
+            ## the experiments, so it is left as-is for reproducibility. The following commented-out
+            ## code could be used for intended DPO behavior, but then to reproduce the paper's qualitative 
+            ## reward findings (where CPC could improve reward), additional hyperparameter tuning may be required,
+            ## for instance increasing DPO learning rate to make it a bit overly aggressively optimized. 
+            ## In addition to this line, the commented-out line "scores_np = np.where(np.isnan(scores_np), np.inf, scores_np)"
+            ## in ../data/synthetic_dataset_formatter.py should also be incorporated.
+
+            # if input_scores[i] is None and (score is not None and score != float("inf")):
+            #     num_decreased_score += 1
+            # elif input_scores[i] is not None and score < input_scores[i]:
+
             if (input_scores[i] is not None and score is not None) and score < input_scores[i]:
                 num_decreased_score += 1
             if score == float("inf"):