Add VAS to TRL

[idanshen]

This PR introduces the implementation of Value Augmented Sampling (VAS), based on the paper Value Augmented Sampling for Language Model Alignment and Personalization. VAS enhances language model generation by incorporating a value function that guides the decoding process.

Key Components:

• VASTrainer: Implements the training algorithm for the Value function.
• VASSampler: Implements the generation process according to the VAS decoding scheme. Uses the trained Value function to adjust the sampling probabilities during inference.

[idanshen]

Hi! Just wanted to check in—I'm here if you have any questions or need any clarification on the PR.

[kashif]

@idanshen so would you mind adding a test as well for the trainer?

[kashif]

@idanshen also documentation would be awesome as well as some results if you have to share?

[idanshen]

Few changes:

• Update VASTrainer code structure to be similar to the new PPOTrainer
• Add test for VASTrainer
• Add documentation
• Example of VASTrainer run - https://api.wandb.ai/links/idanshen/x2ah1rva