[CGPO] Mixture of judges

[gaetanlop]

What does this PR do?

This PR adds the Mixture of judges part of the CGPO paper (https://arxiv.org/pdf/2409.20370). The judges are described in section 4.1.4 and the mixture of judges simply labels a generation as “violated” (0) if it fails any one of the constraint judgments and “satisfied” (1) otherwise.

Related to #2156

Before submitting

• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines.
• Did you write any new necessary tests?

Who can review?

@kashif @lewtun

[gaetanlop]

Still a draft PR Ready

[qgallouedec]

Thanks a lot @gaetanlop. Added some suggestion and open questions

[qgallouedec]

Also, please make sure to run the pre-commits (make precommit)

[gaetanlop]

Hey @qgallouedec, thanks for the review. I have added a gold_answers parameter to the judge function for both the base and MoJs class. Also, I have added a safety_judge and factuality_judge as described in the CGPO paper. They also have rule based judges, but in my opinion they are a little bit too tailored to specific tasks (coding/maths) to be added to the trl library. If you think the factuality and safety judges are also too specific to be in the lib I can remove them from the PR.

For the naming of the judges, let's do BaseBinaryJudge for the base class and AllTrueJudge for the moj following your suggestions?

[qgallouedec]

Hey @qgallouedec, thanks for the review. I have added a gold_answers parameter to the judge function for both the base and MoJs class. Also, I have added a safety_judge and factuality_judge as described in the CGPO paper. They also have rule based judges, but in my opinion they are a little bit too tailored to specific tasks (coding/maths) to be added to the trl library. If you think the factuality and safety judges are also too specific to be in the lib I can remove them from the PR.

Thanks a lot! Nice work!

For the naming of the judges, let's do BaseBinaryJudge for the base class and AllTrueJudge for the moj following your suggestions?

LGTM.

---

WDYT of having generic classes in trl.judges (AllTrueJudge, BinaryJudge etc.) and subclass them in trl.trainer.cgpo_trainer to get SafetyConstraintJudge and FacultyConstraintJudge? If in the future we need these classes elsewhere, we can still move them in trl.judges

[gaetanlop]

Ok, I have removed the FacultyConstraintJudge and the SafetyConstraintJudge from the PR and made the required renamings. Thanks @qgallouedec for the feedback.

[qgallouedec]

maybe "gold_completion" or even "ref_completion" would be more suited, what do you think?

[gaetanlop]

"gold_completions" seems good to me. I have done the update.

[qgallouedec]

This should return a list of int to be consistent with the super class

[gaetanlop]

Indeed!