How to attribute reward to multiple model runs in the same trajectory with PPO

I want to finetune a base model `M` to maximize a reward `R`, when the model is used inside of a more complex system.
Take a simple example of the setting. The trajectory is as follows: sample `prompt_1` from a dataset of prompts, then
```
prompt1 -> M(prompt1) = out_1
out_1 -> F(out_1) = prompt_2
prompt_2 -> M(prompt_2) = out_2
out_2 -> R(out_2) = reward
```

where `F : str -> str` and `R : str -> int` are some methods defined in my code.
Is there a way to do this in the current TRLX framework, preferably online with PPO?
Alternative suggestions are welcome.

Review TRLX's current PPO trajectory and reward-handling entry points to determine whether one trajectory can contain multiple model runs linked to a single reward. Confirm whether the existing online PPO flow supports this; done means documenting the supported approach or defining the missing framework behavior.

Written by the indexing model from the issue text.