Categorical distribution for discrete actions. Supports invalid action masking. :param action_dim: Number of discrete actions

Dataclass for the output of the method policy.evaluate_actions(). This is invoked during training phase for each mini-batch in the rollout bu

Dataclass for the output of the method policy.foward_policy()

Base class for outputs of decoder-only generation models using beam search. Args: sequences (`torch.LongTensor` of shape `(batch_siz

Base class for outputs of encoder-decoder generation models using beam search. Hidden states and attention weights of the decoder (respective

Main Class to compute the CIDEr metric

CIDEr scorer.

Modified PyTorch Categorical distribution with support for invalid action masking. To instantiate, must provide either probs or logits, but

Dataclass for the output of the method policy.get_ref_log_probs()

Main Class to compute the SPICE metric

Dataclass for the output of the method policy.forward_value()

Base class for outputs of decoder-only generation models using beam sample. Args: sequences (`torch.LongTensor` of shape `(batch_siz

Base class for outputs of encoder-decoder generation models using beam sampling. Hidden states and attention weights of the decoder (respecti

Base class for outputs of decoder-only generation models using greedy search. Args: sequences (`torch.LongTensor` of shape `(batch_

Base class for outputs of encoder-decoder generation models using greedy search. Hidden states and attention weights of the decoder (respecti

Bernoulli distribution for multibinary actions. Supports invalid action masking. :param action_dim: Number of binary actions

Dict Rollout buffer used in on-policy algorithms like A2C/PPO. Extends the RolloutBuffer to use dictionary observations It corresponds t

Callback for evaluating an agent. Supports invalid action masking. :param eval_env: The environment used for initialization :param callb

MultiCategorical distribution for multi discrete actions. Supports invalid action masking. :param action_dims: List of sizes of discrete act

A generic trainer for training LMs with onpolicy algorithms from SB3

Official version

Base class for outputs of decoder-only generation models using sampling. Args: sequences (`torch.LongTensor` of shape `(batch_size*

Base class for outputs of encoder-decoder generation models using sampling. Hidden states and attention weights of the decoder (respectively

A supervised trainer to train LMs (causal and seq2seq) on text generation tasks (wrapper on HF trainer)

Advantage Actor Critic (A2C) Paper: https://arxiv.org/abs/1602.01783 Code: This implementation borrows code from https://github.com/ikos

Official version

A base class for all the environments

Placeholder for observation data class

Computes rewards for several instances at once

A class containing all functions for auto-regressive text generation, to be used as a mixin in [`PreTrainedModel`]. The class exposes [`~gen

IMDB Dataset for sentiment continuation task

IMDB Dataset in seq2seq format to train supervised generator

CNN policy class for actor-critic algorithms (has both policy and value prediction). Used by A2C, PPO and the likes. :param observation_

Policy class for actor-critic algorithms (has both policy and value prediction). Used by A2C, PPO and the likes. :param observation_spac

MultiInputActorClass policy class for actor-critic algorithms (has both policy and value prediction). Used by A2C, PPO and the likes. :p

Rollout buffer that also stores the invalid action masks associated with each observation. :param buffer_size: Max number of element in the