NVIDIA-NeMo / NVIDIA-NeMo/RL

training on transitions

Open
#1,326 1 comment 0 reactions 0 assignees View on GitHub
enhancement r0.6.0 x-futurehouse
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Is your feature request related to a problem? Please describe.**
Option to make groups in GRPO out of individual steps rather than full trajectories when using multi-step environments.
Rather than N trajectories per group, N steps per group, constructed from 1 (or more?) trajectories.

**Describe alternatives you've considered**
Nemo RL supports step-level rewards but they are aggregated and applied at a trajectory level only.

**Additional context**
Useful for agent model training

**Example**
**Sliding puzzle:**
For each task N rollouts are generated to make a group, each with multiple steps/turns.

Group 1: Task 1
(User, Assistant Action, Observation, Action, Observation ..., Reward) x N

Group 2: Task 2
(U,A,O,A,O,...,R) x N

...

**Training on transitions / grouping individual steps would be like:**
Group 1: Task 1 (rollout 1)
(U,A,O,R)
(U,A,O,A,O,R)
(U,A,O,A,O,A,O,R)
...

Group 2: Task 1 (rollout 2)
(U,A,O,R)
(U,A,O,A,O,R)
(U,A,O,A,O,A,O,R)
...

or maybe each element in group is just (Instruction,Latest Observation,Action,Reward) dropping previous message history.

I will try to open a PR

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.