gflownet.policy.simple_tree

Classes

Backbone	GNN backbone: a stack of GNN layers that can be used for processing graphs.
Head	GNN head: a combination of a pooling layer and a stack of linear layers,
SimpleTreeModel	Combination of backbone and head models. In forward, it converts the states
SimpleTreePolicy	Policy wrapper using SimpleTreeModel as the policy model.

Module Contents

class gflownet.policy.simple_tree.Backbone(input_dim, n_layers=3, hidden_dim=128, layer='GINEConv', activation='LeakyReLU')

Bases: torch.nn.Module

GNN backbone: a stack of GNN layers that can be used for processing graphs.

Parameters:

• input_dim (int)

• n_layers (int)

• hidden_dim (int)

• layer (str)

• activation (str)

hidden_dim = 128

model

forward(data)

Parameters:

data (torch_geometric.data.Data)

Return type:

torch.Tensor

class gflownet.policy.simple_tree.Head(backbone, out_dim, n_layers=2, hidden_dim=256, activation='LeakyReLU')

Bases: torch.nn.Module

GNN head: a combination of a pooling layer and a stack of linear layers, that takes the input graphs processed by the provided backbone and outputs out_dim values.

This is a naive variant of a Tree policy model, that in particular doesn’t do any node-level prediction, but instead predicts the whole policy output all at once.

Parameters:

• backbone (torch.nn.Module)

• out_dim (int)

• n_layers (int)

• hidden_dim (int)

• activation (str)

backbone

head

forward(data)

Parameters:

data (torch_geometric.data.Data)

Return type:

(torch.Tensor, torch.Tensor)

class gflownet.policy.simple_tree.SimpleTreeModel(n_features, policy_output_dim, base=None, backbone_args=None, head_args=None)

Bases: torch.nn.Module

Combination of backbone and head models. In forward, it converts the states to a batch of graphs, and passes them through both models.

Parameters:

• n_features (int)

• policy_output_dim (int)

• base (Optional[SimpleTreePolicy])

• backbone_args (Optional[dict])

• head_args (Optional[dict])

n_features

policy_output_dim

model

forward(x)

class gflownet.policy.simple_tree.SimpleTreePolicy(config, env, device, float_precision, base=None)

Bases: gflownet.policy.base.Policy

Policy wrapper using SimpleTreeModel as the policy model.

backbone_args

head_args

n_features

policy_output_dim

is_model = True

parse_config(config)

instantiate()

__call__(states)