NVIDIA-NeMo / NVIDIA-NeMo/RL

Config-passing custom reward functions

Open
#1,211 2 comments 0 reactions 0 assignees View on GitHub
community-request enhancement external r0.6.0 waiting-on-customer x-atlassian
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.