GO-CRE

Guided Optimization of Cis-Regulatory Elements — a reinforcement-learning framework for in silico directed evolution of synthetic cell-type-specific enhancers, built on the HybriDNA hybrid Transformer-Mamba2 long-range DNA language model.

This is the companion code for "Reinforcement learning guided directed evolution of synthetic cis-regulatory elements" (Ma, Bu, Liu, et al., 2026). GO-CRE drives PPO on a prompt-conditioned HybriDNA-300M policy with a Malinois (CODA) regulatory-activity reward to design 200-bp synthetic enhancers with target-cell-type-selective activity in K562, HepG2, and SK-N-SH; the trajectory of optimization is then deconvolved into an interpretable sequence-grammar landscape.

---

Layout

GO-CRE/
├── hybridna/      foundation model (Mamba2 + GQA hybrid, RC-echo head)
├── gocre/         paper training framework: SFT + PPO + Malinois reward
├── analysis/      bioinformatic analysis pipeline (k-mer/NMF/PCA, motifs, MPRA library design)
├── examples/      HybriDNA usage demos: embedding, generation, fine-tuning
├── scripts/       driver scripts (run_ppo.sh)
├── docs/          paper Methods reference
├── supplementary/ paper supplementary materials: generated sequences, MPRA libraries, video
├── pyproject.toml · requirements.txt   pip install -e .
└── LICENSE        MIT

---

Install

# Python ≥ 3.11, CUDA ≥ 12.4
pip install -e .                  # installs hybridna and gocre as editable packages
# or core deps only:
pip install -r requirements.txt

Additional analysis dependencies (R + Python) are listed in analysis/README.md.

---

Quick start — train a cell-type-specific CRE generator

1. SFT a HybriDNA-300M on cell-type-prompt-conditioned MPRA data

python -m gocre.sft_train \
    --data data/cre_specific_200bp.jsonl \
    --base-model Mishamq/HybriDNA-300M \
    --output-dir checkpoints/hybridna-cre-sft-300m \
    --num-epochs 3 --batch-size 32 --learning-rate 5e-5
# or after `pip install -e .`:
gocre-sft --data data/cre_specific_200bp.jsonl --output-dir checkpoints/hybridna-cre-sft-300m

The paper's MPRA filter ("high activity in exactly one cell type, low in the other two", from Malinois-scored data) should be applied upstream; this script trusts the labels in cell_type.

2. PPO with the Malinois (CODA) reward, per cell type

bash scripts/run_ppo.sh --cell K562    # erythroid
bash scripts/run_ppo.sh --cell HepG2   # hepatocyte (polyG / G-C guard on)
bash scripts/run_ppo.sh --cell SKNSH   # neuroblastoma (A/T guard on)

The launcher resolves checkpoint / BODA / Malinois locations from env vars (with repo-relative defaults), prints the resolved settings, and writes everything under <repo>/runs/ppo_<cell>_<timestamp>/.

The only differences between cells are the prompt token and the homopolymer guardrail:

Cell	Prompt	A/T guard	G/C guard (polyG)
K562	100	off	off
HepG2	010	off	on (weight=0.06, threshold=4) — paper-validated polyG penalty
SKNSH	001	on (weight=0.06, threshold=4) — paper-validated A/T guard	off

Both branches share the same soft-penalty form: weight · (max(0, max_X_run − threshold) + max(0, max_Y_run − threshold)) is subtracted from the reward, where (X, Y) is (A, T) or (G, C). See compute_low_complexity_penalty in gocre/ppo_trainer.py and the deeper writeup in docs/method.md.

3. Calling the trainer programmatically

from gocre import PPOConfig, PPOTrainer

config = PPOConfig(
    prompts=["010"],                # HepG2
    num_iterations=1000,
    model_path="/path/to/sft/ckpt",
    output_dir="./runs/hepg2_test",
    use_bf16=True,
    use_minigap=True,
    max_gc_run_penalty_weight=0.06, # polyG guard for HepG2
    max_gc_run_threshold=4,
)
trainer = PPOTrainer(config)
trainer.train()

Prerequisites for the PPO step

1. HybriDNA-300M SFT checkpoint — produced by step 1 above. Default location: <repo>/checkpoints/hybridna-cre-sft-300m.
2. CODA / boda2 — clone https://github.com/sjgosai/boda2 so it imports as boda (the trainer does import boda at runtime). Default location: <repo>/boda2.
3. Malinois artifacts — download malinois_artifacts__20211113_021200__287348.tar.gz from gs://tewhey-public-data/CODA_resources/. Default location: <repo>/resources/. The trainer falls back to remote download if no local copy is found.

Override paths via env vars (PPO_MODEL_PATH, BODA_ROOT, MALINOIS_ARTIFACT_PATH, PPO_OUTPUT_DIR) or pass extra --flag arguments to run_ppo.sh — they're forwarded to the trainer.

---

HybriDNA model

HybriDNA is the underlying DNA foundation model — it interleaves Mamba-2 state-space layers with Grouped Query Attention at a 7:1 ratio, supports sequences up to 131,074 bp, and ships a reverse-complement (RC) echo embedding for bidirectional discriminative tasks.

Model	Parameters	Hidden	Layers	Max length	Hugging Face
HybriDNA-300M	300M	1024	24	131,074 bp	Mishamq/HybriDNA-300M
HybriDNA-3B	3B	4096	16	131,074 bp	Mishamq/HybriDNA-3B
HybriDNA-7B	7B	4096	32	131,074 bp	Mishamq/HybriDNA-7B

The same modeling code is mirrored to the Hub for the published checkpoints, so loading with trust_remote_code=True works out of the box. After pip install -e . you can also drop trust_remote_code=True and use the standard Auto* classes directly.

Embedding extraction

from transformers import AutoTokenizer, AutoModel
import torch

model_name = "Mishamq/HybriDNA-300M"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)

inputs = tokenizer("ACGTACGTACGTACGT", return_tensors="pt")
with torch.no_grad():
    embeddings = model(**inputs).last_hidden_state            # (B, L, H)
sequence_embedding = embeddings.mean(dim=1)                   # (B, H)

Fine-tuning with RC echo embeddings

HybriDNAForSequenceClassificationEcho concatenates the input with itself (or its reverse complement), runs the autoregressive model over the doubled sequence, and pools the second half — giving a decoder-only model bidirectional context. The RCE variant is the one reported in the paper.

from hybridna import HybriDNAForSequenceClassificationEcho, HybriDNATokenizer
import torch

ckpt = "Mishamq/HybriDNA-300M"
tokenizer = HybriDNATokenizer.from_pretrained(ckpt)
model = HybriDNAForSequenceClassificationEcho.from_pretrained(
    ckpt, num_labels=2, torch_dtype=torch.bfloat16,
)

def rc(s): return s.translate(str.maketrans("ACGT", "TGCA"))[::-1]
seq = "ACGTACGTACGT"
batch = tokenizer(seq + rc(seq), return_tensors="pt")  # RCE
logits = model(**batch).logits

See examples/finetune_echo.py for a runnable script using HF Trainer.

---

Supplementary materials

The supplementary/ directory bundles the data released with the paper:

• Generated sequences — 9000 GO-CRE-designed 200-bp enhancers per target cell type (HepG2, K562, SK-N-SH), with Malinois-predicted activity scores in the FASTA headers.
• MPRA libraries — Library A (500 reference sequences + reverse complements used for activity modeling) and Library B HepG2 validation set (1000 sequences: generated designs + STARR-derived positives + negative controls).
• Supplementary Video 1 — supplement_video1.mp4.

See supplementary/README.md for file-by-file details (header format, counts, source).

---

Citation

If you use GO-CRE, please cite the directed-evolution paper:

@article{ma2026gocre,
  title   = {Reinforcement learning guided directed evolution of synthetic
             cis-regulatory elements},
  author  = {Ma, Mingqian and Bu, Wanjuan and Liu, Guoqing and others},
  year    = {2026},
}

If you also use the underlying HybriDNA model, please cite:

@article{ma2025hybridna,
  title  = {HybriDNA: A Hybrid Transformer-Mamba2 Long-Range DNA Language Model},
  author = {Ma, Mingqian and Liu, Guoqing and Cao, Chuan and Deng, Pan and
            Dao, Tri and Gu, Albert and Jin, Peiran and Yang, Zhao and
            Xia, Yingce and Luo, Renqian and others},
  journal = {arXiv preprint arXiv:2502.10807},
  year    = {2025}
}

---

License

MIT — see LICENSE.