gfn.modules

Module Contents

Classes

GFNModule	Abstract Base Class for all functions/approximators/estimators used.
NeuralNet	Abstract Base Class for all functions/approximators/estimators used.
Tabular	Abstract Base Class for all functions/approximators/estimators used.
Uniform	Use this module for uniform policies for example. This is because logits = 0 is equivalent to uniform policy
ZeroGFNModule	Abstract Base Class for all functions/approximators/estimators used.

Attributes

InputTensor
OutputTensor

class gfn.modules.GFNModule

Bases: abc.ABC

Abstract Base Class for all functions/approximators/estimators used. Each module takes a preprocessed tensor as input, and outputs a tensor of logits, or log flows. The input dimension of the module (e.g. Neural network), is deduced from the environment’s preprocessor’s output dimension

abstract property output_dim: int

Output dimension of the module

Return type

int

abstract __call__(preprocessed_states)

Parameters

preprocessed_states (InputTensor) –

Return type

OutputTensor

abstract load_state_dict(state_dict)

Parameters

state_dict (dict) –

named_parameters()

Returns a dictionary of all (learnable) parameters of the module. Not needed for NeuralNet modules, given that those inherit this function from nn.Module

Return type

dict

gfn.modules.InputTensor

class gfn.modules.NeuralNet(input_dim, output_dim, hidden_dim=256, n_hidden_layers=2, activation_fn='relu', torso=None)

Bases: torch.nn.Module, GFNModule

Abstract Base Class for all functions/approximators/estimators used. Each module takes a preprocessed tensor as input, and outputs a tensor of logits, or log flows. The input dimension of the module (e.g. Neural network), is deduced from the environment’s preprocessor’s output dimension

Parameters

• input_dim (int) –

• output_dim (int) –

• hidden_dim (Optional[int]) –

• n_hidden_layers (Optional[int]) –

• activation_fn (Optional[Literal[relu, tanh]]) –

• torso (Optional[torch.nn.Module]) –

property output_dim: int

Output dimension of the module

Return type

int

forward(preprocessed_states)

Parameters

preprocessed_states (InputTensor) –

Return type

OutputTensor

gfn.modules.OutputTensor

class gfn.modules.Tabular(n_states, output_dim)

Bases: GFNModule

Abstract Base Class for all functions/approximators/estimators used. Each module takes a preprocessed tensor as input, and outputs a tensor of logits, or log flows. The input dimension of the module (e.g. Neural network), is deduced from the environment’s preprocessor’s output dimension

Parameters

• n_states (int) –

• output_dim (int) –

property output_dim: int

Output dimension of the module

Return type

int

__call__(preprocessed_states)

Parameters

preprocessed_states (InputTensor) –

Return type

OutputTensor

load_state_dict(state_dict)

Parameters

state_dict (dict) –

named_parameters()

Returns a dictionary of all (learnable) parameters of the module. Not needed for NeuralNet modules, given that those inherit this function from nn.Module

Return type

dict

class gfn.modules.Uniform(output_dim)

Bases: ZeroGFNModule

Use this module for uniform policies for example. This is because logits = 0 is equivalent to uniform policy

Parameters

output_dim (int) –

class gfn.modules.ZeroGFNModule(output_dim)

Bases: GFNModule

Abstract Base Class for all functions/approximators/estimators used. Each module takes a preprocessed tensor as input, and outputs a tensor of logits, or log flows. The input dimension of the module (e.g. Neural network), is deduced from the environment’s preprocessor’s output dimension

Parameters

output_dim (int) –

property output_dim: int

Output dimension of the module

Return type

int

__call__(preprocessed_states)

Parameters

preprocessed_states (InputTensor) –

Return type

OutputTensor

load_state_dict(state_dict)

Parameters

state_dict (dict) –