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.

[dpaleka]

As of now, TRLX supports only RL setups where all “actions” to attribute the reward to are done before the reward function is called.

@dpaleka, isn’t this already the case in your very first pseudocode snippet? R is only called after both M calls, not in between, right?

I miswrote; what is true is that TRLX assumes the “actions” are just a single model.generate call.

[paulbricman]

@paulbricman are you on the discord? https://discord.gg/canadagoose

Yes, just @paul on that server.

Hey! I am open to rectifying this, I am just at capacity right now and I don’t think we have the engineering manpower for this at the moment. @paulbricman @dpaleka if you two would be interested in implemented, I’d be happy to assign you and then review it.

I’d be excited to help implement it, but I’m skeptical about whether I understand PPO well enough and whether I’m familiar enough with the trlx codebase to do it. I might be able to make a contribution if both @dpaleka and I would work on this?

Also thanks for the super quick reply!