LMHeadActorValueOperator

class torchrl.modules.tensordict_module.LMHeadActorValueOperator(*args, **kwargs)

Builds an Actor-Value operator from an huggingface-like *LMHeadModel.

This method:

• takes as input an huggingface-like *LMHeadModel

• extracts the final linear layer uses it as a base layer of the actor_head and adds the sampling layer

• uses the common transformer as common model

• adds a linear critic

Parameters:

base_model – a torch model composed by a .transformer model and .lm_head linear layer