gflownet.losses.trajectorybalance
=================================

.. py:module:: gflownet.losses.trajectorybalance

.. autoapi-nested-parse::

   Trajectory Balance loss or objective for training GFlowNets.

   The Trajectory Balance (TB) loss or objective was defined by Malkin et al. (2022):

       .. _a link: https://arxiv.org/abs/2201.13259



Classes
-------

.. autoapisummary::

   gflownet.losses.trajectorybalance.TrajectoryBalance


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

.. py:class:: TrajectoryBalance(**kwargs)

   Bases: :py:obj:`gflownet.losses.base.BaseLoss`


   Initialization method for the Trajectory Balance loss class.

   .. attribute:: name

      The name of the loss or objective function: Trajectory Balance

      :type: str

   .. attribute:: acronym

      The acronym of the loss or objective function: TB

      :type: str

   .. attribute:: id

      The identifier of the loss or objective function: trajectorybalance

      :type: str


   .. py:attribute:: name
      :value: 'Trajectory Balance'



   .. py:attribute:: acronym
      :value: 'TB'



   .. py:attribute:: id
      :value: 'trajectorybalance'



   .. py:method:: requires_backward_policy()

      Returns True if the loss function requires a backward policy.

      The Trajectory Balance loss does require a backward policy model, hence True is
      returned.

      :returns: *True*



   .. py:method:: requires_state_flow_model()

      Returns True if the loss function requires a state flow model.

      The Trajectory Balance loss does not require a state flow model, hence False is
      returned.

      :returns: *False*



   .. py:method:: is_defined_for_continuous()

      Returns True if the loss function is well defined for continuous GFlowNets,
      that is continuous environments, or False otherwise.

      The Trajectory Balance loss is well defined for continuous GFlowNets, therefore
      this method returns True.

      :returns: *True*



   .. py:method:: compute_losses_of_batch(batch)

      Computes the Trajectory Balance loss for each trajectory of the input batch.

      The Trajectory Balance (TB) loss or objective is computed in this method as is
      defined in Equation 14 of Malkin et al. (2022).

      .. _a link: https://arxiv.org/abs/2201.13259

      :param batch: A batch of trajectories.
      :type batch: Batch

      :returns: *tensor* -- The loss of each trajectory in the batch.



   .. py:method:: aggregate_losses_of_batch(losses, batch)

      Aggregates the losses computed from a batch to obtain the overall average loss.

      The result is returned as a dictionary with the following items:
      - 'all': Overall average loss

      :param losses: The loss of each trajectory in the batch.
      :type losses: tensor
      :param batch: A batch of trajectories.
      :type batch: Batch

      :returns: **loss_dict** (*dict*) -- A dictionary of loss aggregations.