Lessons from Designing Rewards for Crystal Generation

It has long been known that constructing a reward function on top of an environment, such that a desired objective is learned, is a non-trivial task. Among the many problems that arise, reward hacking stands out the most due to its catastrophic results and its tendency to appear where you least expect it.

Why is it so hard? In my experience, the root cause is that you start with an objective that is inherently difficult to define precisely: "What is a clean room? What is a good poem? What is a stable crystal?" Beyond that, the limitations of your design are impossible to predict in advance; they only surface after training has already revealed the gaps. This is compounded by the model itself: when a model is not capable enough to reliably satisfy a well-defined reward, the reward designer is forced to make strategic tradeoffs between training longer, running faster iterations with tweaked rewards, or restructuring the reward entirely.

What follows are five concrete lessons I learned while designing the reward function for crystal generation using GRPO on 0.5–8B parameter models. Some of these lessons emerged quickly; others only became clear after watching hundreds of failed rollouts. All of them are grounded in specific failures from the crystal generation task.

Why This Matters for Crystal Generation Specifically

This is really an issue of model capability. If the models I was using were sufficiently powerful, a simple binary reward (is the crystal stable or not?) would be enough. But at the scales tested (0.5–8B parameters), they simply are not. Even the original DeepSeek paper included a lightweight format reward term rather than discarding any response that lacked proper formatting, because a model that can't yet produce valid outputs still needs a gradient signal to learn from.

Crystal structure generation adds further difficulty: the problem is highly underdetermined (many valid CIFs exist for a given composition), the evaluation oracle (CHGNet) has its own noise floor, and there is no natural language description of what a "slightly better" crystal looks like. Every lesson below was hard-won from the specific constraints of this task.

Lesson 1: Rubric-Style Rewards Beat Gated Rewards

My first reward function was sequential: it checked format, then composition, then bond lengths, then forces. If any check failed, the function returned 0 immediately. The logic seemed defensible: why reward a crystal that doesn't even have the right composition? But the consequence was that a large fraction of the most valuable training examples were silently discarded.

Here is a concrete example. Early in training, the model occasionally generated a CIF file with excellent bond geometry and low forces, a structurally coherent proposal, but with a composition that was slightly off (e.g., NaCl₁.₅ instead of NaCl). Under the gated reward, this rollout earned exactly 0. Under a rubric reward, it would have earned partial credit for bond lengths and forces, providing a real gradient signal on two of the four sub-goals.

def reward(cif, target):
    if not parse(cif):
        return 0.0
    if not composition_ok(cif, target):
        return 0.0       # ← silent zero
    if not bonds_ok(cif):
        return 0.0
    if not forces_ok(cif):
        return 0.0
    return 1.0

def reward(cif, target):
    r = 0.0
    r += 0.10 * format_score(cif)
    r += 0.15 * composition_score(cif, target)
    r += 0.15 * bond_score(cif)
    r += 0.24 * forces_score(cif)
    r += 0.60 * energy_score(cif)
    return r

The deeper problem with gated rewards is that they enforce a strict learning order: the model must master step N before it receives any signal about step N+1. In practice, this means the model can spend many steps stuck on composition while its capacity to learn bond geometry, which it might naturally acquire in parallel, goes entirely unused. Switching to rubric-style scoring accelerated convergence across all sub-goals simultaneously.

The rubric approach does introduce the multi-term complexity from Lesson 3, but the tradeoff is worth it. Dense gradient coverage is more valuable than the simplicity of a clean sequential gate.

Lesson 2: Avoid Negative Rewards

Early in the project, I ran into a common failure mode: the model was generating text that resembled a CIF file but would fail parsing for superficial reasons: content outside the answer tags, trailing commentary mixed into the structure, or minor syntax errors. My first instinct was to add a negative reward whenever the content inside <|answer_start|>...<|answer_end|> was not parseable by pymatgen's CifParser.

This was a mistake.

The reward landscape at that point in training looked roughly like this:

The model found the optimal policy immediately: output empty tags and do nothing inside. From its perspective, this was rational. Just using the answer tags earned the format reward (+0.10) regardless of what was inside. Attempting to include a CIF but getting it wrong only netted +0.05, strictly worse than doing nothing. And the probability of generating a valid CIF from scratch was too low early in training to shift the expected value. So the policy quickly converged to: earn the safe 0.10, then stop.

Entropy collapsed almost immediately. The model had found a local optimum, and the very low probability of any individual rollout producing a valid CIF was never high enough to push it out. This is the classic hill-climbing trap: a small guaranteed reward for the safe action dominates when the risky alternative returns less.

The fix was straightforward: remove the negative reward entirely. Now, producing empty tags still earns +0.10, but producing an invalid CIF earns 0.0 instead of being penalized. This keeps exploration alive. The model no longer has a reason to prefer empty tags over a genuine attempt. The rare valid CIF, when it appears, reinforces that behavior without the gradient being suppressed by penalties elsewhere.

Lesson 3: The More Terms, the More Unexpected the Outcome

To see why, consider a reward with three independent binary judges (each returning 0 or 1). The set of achievable total rewards is {0, 1, 2, 3}. Now suppose one judge is weighted at 0.5× instead of 1.0×. The achievable set expands to {0, 0.5, 1, 1.5, 2, 2.5}: only three judges, but 50% more distinct reward levels. Crucially, the mapping from policy behavior to reward is now less unique:

With equal weights, Policies A, B, and C are indistinguishable (all score 2). Discount Judge 1 to 0.5× and Policy C jumps to the top (2.0) while Policies A and B both drop to 1.5. Policy C never passes Judge 1 at all, yet it wins because it doesn't need to. This is the core mechanic behind composition hacking in the crystal generation task: the composition rubric is Judge 1.

The Crystal Generation Case: Composition Hacking

In my first multi-term iteration, I used four rubrics: format, composition, bond lengths, and forces. Composition was included intentionally. Free generation without a compositional constraint would inadvertently reward simpler crystals, since binary metals and symmetric oxides are far easier to generate than ternary or quaternary systems.

Training proceeded well at first. Then, around step 1000, I noticed the aggregate reward kept climbing while the composition sub-score steadily declined. The model had discovered a simple exploit: ignore the target composition and generate an easier, structurally simpler crystal instead. A binary metallic oxide, for example, is far more likely to achieve low formation energy and reasonable bond lengths than a ternary chalcogenide. Two rubrics improving more than compensated for one rubric getting worse.

The fix, for smaller models where this was most acute, was to temporarily increase the composition weight to force the policy back into the compositionally constrained region, then gradually rebalance as the model became capable of satisfying all rubrics simultaneously. At larger scales (30B and 230B parameters), this problem was much less severe: those models had enough capacity to satisfy composition and stability constraints at the same time, so inflated composition weights were not needed.

Lesson 4: Continuous Rewards Outperform Binary Ones

For the formation energy rubric, the materials science community uses a canonical threshold: structures with E_f < 0.1 eV/atom are considered thermodynamically stable and likely synthesizable. The natural instinct is to implement this as a step function: above threshold → 0, below threshold → 1.

The problem is that formation energy is a rich, continuous signal. A structure with E_f = 0.5 eV/atom is meaningfully better than one with E_f = 3.2 eV/atom: it is closer to stability, its bond geometry is more coherent, its forces are lower. But under a step reward, they receive identical scores (0). The policy has no reason to prefer one over the other.

Compare the three available reward strategies:

def formation_energy_discrete(energy, threshold=0.1):
    return 1.0 if energy <= threshold else 0.0

def formation_energy_linear(energy, threshold=0.1, max_penalty=2.0):
    if energy <= threshold:
        return 1.0
    elif energy >= threshold + max_penalty:
        return 0.0
    else:
        return 1.0 - (energy - threshold) / max_penalty

def formation_energy_exponential(energy, threshold=0.1, scale=1.0):
    if energy <= threshold:
        return 1.0
    excess = energy - threshold
    return np.exp(-excess / scale)

The linear and exponential variants both provide useful gradient signal for structures that are not yet stable, guiding the policy to incrementally reduce formation energy rather than treating all sub-threshold structures as equally worthless. In practice, the linear variant is preferable for early training because the flat gradient near the threshold doesn't create a "cliff" that destabilizes learning once the model is close to passing.

There are cases where binary rewards are appropriate: when the task is genuinely all-or-nothing and there is no meaningful partial credit (a correct classification, a flawless backflip). For crystal generation, this is not the case. The metric is continuous, the DFT oracle is noisy, and dense gradient coverage is more important than strict threshold semantics.

Lesson 5: Always Think About the Path of Least Resistance

This is less of a standalone lesson and more of a meta-rule that encompasses the others. Every reward hacking failure I encountered was, in retrospect, a failure to predict the path of least resistance in advance.

The empty-tag exploit (Lesson 2) was the path of least resistance when negative rewards made exploration costly. Composition hacking (Lesson 3) was the path of least resistance when simpler crystals were easier to generate and two other rubrics were sufficient to compensate. The gated reward (Lesson 1) was the path of least resistance for the training loop itself; it was easier to implement but wasteful in practice.

A useful exercise before training: for each reward term, write down the policy that maximizes that term in isolation while actively failing all others. Then ask: does the total reward still incentivize that degenerate policy? If yes, rebalance. The composition example from Lesson 3 is a perfect illustration: a model that generates whatever crystal it finds easiest, regardless of the target formula, will consistently outscore a model that tries to honor the composition constraint but generates slightly imperfect structures.

How These Lessons Shaped the Final Reward

The reward function that trained the full system incorporates all five lessons. No negative rewards. Independent rubric scoring. Continuous rewards for energy and forces where possible. Weights tuned to close off known exploitation paths.

None of these weights emerged from a theoretic argument alone. They were each tested against the pre-training audit: which degenerate policies does this weighting allow, and does the full reward still penalize them adequately? The 60% formation energy weight, in particular, is high enough to make purely compositional hacking non-viable; the model needs to hit near-stable energies, not just the right stoichiometry, to score well.