Config-passing custom reward functions
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
It would be good to have a low touchpoint way to use custom reward functions without directly modifying the library sourcecode, or creating a new training recipe. Note this would apply only to the VLMEnvironment and any training recipes using it.
I think we could replace the [if/else logic](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/environments/vlm_environment.py#L71) with a dict mapper (like `task_to_env` but `name_to_reward`). Then the user could pass into the config which is parsed to extract additional functions and merge them into that mapper. e.g.:
* a path to another dict mapper - either a script or module on the import path - and the name of a dict variable containing the mapper, or
* a list of path/name pairs
Note that I'm undecided whether there's any more elegance to supporting scripts or modules directly. Supporting a script should be easy because it can always `*` import any rewards from modules that the user desires. This way a user with custom reward requirements only needs one simple file colocating their reward functions.
[`verl` has a similar setup](https://github.com/volcengine/verl/blob/377bbb84f068e7ca0205ecfaf7841878c2a27b66/verl/trainer/ppo/reward.py#L42) where a python script containing functions in its scope can be imported. I believe they accept only one function. Ours can be more flexible, taking a dict or list, since the VLM environment itself handles reward function selection logic.
I am happy to contribute this. I just want to confirm whether this would be considered valuable, and whether there is any preference as to the pattern to follow.
Contributor guide
Assessment
This issue has not been assessed yet.