Containers

ModuleExtended

class torchrl.nn.ModuleExtended[source]

Bases: torch.nn.modules.module.Module

A torch module with added functionalities.

SequentialExtended

class torchrl.nn.SequentialExtended(*args, **kwargs)[source]

Bases: torchrl.nn.container.ModuleExtended

A torch sequential module with added functionalities.

FlattenLinear

class torchrl.nn.FlattenLinear(in_features, out_features, **kwargs)[source]

Bases: torch.nn.modules.linear.Linear

Flatten the input and apply a linear layer.

Parameters:
  • in_features (list) – Size of each input sample.
  • out_features (list) – Size of each output sample.

ActionLinear

class torchrl.nn.ActionLinear(in_features, action_info, **kwargs)[source]

Bases: torch.nn.modules.module.Module

A linear layer that automatically calculates the output shape based on the action_info.

Parameters:
  • in_features (list) – Size of each input sample
  • action_info (dict) – Dict containing information about the environment actions (e.g. shape).