tutorials.examples.train_bayesian_structure
===========================================

.. py:module:: tutorials.examples.train_bayesian_structure

.. autoapi-nested-parse::

   The goal of this script is to reproduce the results of DAG-GFlowNet for
   Bayesian structure learning (Deleu et al., 2022) using the GraphEnv.

   Specifically, we consider a randomly generated (under the Erdős-Rényi model) linear-Gaussian
   Bayesian network over `num_nodes` nodes. We generate 100 datapoints from it, and use them to
   calculate the BGe score. The GFlowNet is learned to generate directed acyclic graphs (DAGs)
   proportionally to their BGe score, using the modified DB loss.

   Some expected results on the Erdős-Rényi model with 5 nodes and 5 edges:

   (On-policy training)
   python train_bayesian_structure.py     --seed 0     --no_buffer     --max_epsilon 0.0     --min_epsilon 0.0
   >> Expected SHD: 6.99
   >> Expected edges: 10.00
   >> ROC-AUC: 0.68
   >> Jensen-Shannon divergence: 0.36

   (Off-policy training with epsilon-noisy exploration, which is the default)
   python train_bayesian_structure.py --seed 0
   >> Expected SHD: 6.31
   >> Expected edges: 8.80
   >> ROC-AUC: 0.78
   >> Jensen-Shannon divergence: 0.02



Attributes
----------

.. autoapisummary::

   tutorials.examples.train_bayesian_structure.parser


Classes
-------

.. autoapisummary::

   tutorials.examples.train_bayesian_structure.DAGEdgeActionGNN
   tutorials.examples.train_bayesian_structure.DAGEdgeActionGNNv2
   tutorials.examples.train_bayesian_structure.DAGEdgeActionMLP


Functions
---------

.. autoapisummary::

   tutorials.examples.train_bayesian_structure.main


Module Contents
---------------

.. py:class:: DAGEdgeActionGNN(n_nodes, num_node_classes, num_edge_classes, num_conv_layers = 1, embedding_dim = 128, is_backward = False)

   Bases: :py:obj:`gfn.utils.modules.GraphActionGNN`


   Simple GNN-based edge action module

   :param n_nodes: The number of nodes in the graph.
   :param num_edge_classes: The number of edge classes.
   :param num_conv_layers: The number of GNN layers.
   :param embedding_dim: The dimension of embeddings.
   :param is_backward: Whether the module is used for backward action prediction.


   .. py:property:: edges_dim
      :type: int



   .. py:method:: forward(states_tensor)


   .. py:attribute:: n_nodes


.. py:class:: DAGEdgeActionGNNv2(n_nodes, num_node_classes, num_edge_classes, num_conv_layers = 2, embedding_dim = 128, num_heads = 4, is_backward = False)

   Bases: :py:obj:`torch.nn.Module`


   GNN-based edge action module, adapted from the implementation of
   https://github.com/GFNOrg/GFN_vs_HVI/blob/master/dags/dag_gflownet/nets/gnn/gflownet.py

   :param n_nodes: The number of nodes in the graph.
   :param num_edge_classes: The number of edge classes.
   :param num_conv_layers: The number of GNN layers.
   :param embedding_dim: The dimension of embeddings.
   :param num_heads: The number of attention heads.
   :param is_backward: Whether the module is used for backward action prediction.


   .. py:attribute:: _input_dim
      :value: 1



   .. py:attribute:: _n_nodes


   .. py:attribute:: _output_dim


   .. py:attribute:: attention


   .. py:attribute:: edge_embedding


   .. py:property:: edges_dim
      :type: int



   .. py:attribute:: embedding_dim
      :value: 128



   .. py:method:: forward(states_tensor)


   .. py:attribute:: graph_network


   .. py:property:: input_dim


   .. py:attribute:: is_backward
      :value: False



   .. py:property:: n_nodes


   .. py:attribute:: node_embedding


   .. py:attribute:: num_edge_classes


   .. py:attribute:: num_node_classes


   .. py:property:: output_dim
      :type: int



   .. py:attribute:: projection


   .. py:attribute:: receivers_mlp


   .. py:attribute:: senders_mlp


   .. py:attribute:: stop_mlp


   .. py:attribute:: temperature


.. py:class:: DAGEdgeActionMLP(n_nodes, num_node_classes, num_edge_classes, n_hidden_layers = 3, n_hidden_layers_exit = 2, embedding_dim = 128, is_backward = False)

   Bases: :py:obj:`gfn.utils.modules.GraphEdgeActionMLP`


   Network that processes flattened adjacency matrices to predict graph actions.

   Unlike the GNN-based GraphActionGNN, this module uses standard MLPs to process
   the entire adjacency matrix as a flattened vector. This approach:

   1. Can directly process global graph structure without message passing.
   2. May be more effective for small graphs where global patterns are important.
   3. Does not require complex graph neural network operations.

   The module architecture consists of:
   - An MLP to process the flattened adjacency matrix into an embedding.
   - An edge MLP that predicts logits for each possible edge action.
   - An exit MLP that predicts a logit for the exit action.

   :param n_nodes: Number of nodes in the graph.
   :param directed: Whether the graph is directed or undirected.
   :param n_hidden_layers: Number of hidden layers in the MLP for the edge actions.
   :param n_hidden_layers_exit: Number of hidden layers in the MLP for the exit action.
   :param embedding_dim: Dimension of internal embeddings.
   :param is_backward: Whether this is a backward policy.


   .. py:property:: edges_dim
      :type: int



.. py:function:: main(args)

.. py:data:: parser