# -*- coding: utf-8 -*-
from setuptools import setup

package_dir = \
{'': 'src'}

packages = \
['gfn',
 'gfn.containers',
 'gfn.envs',
 'gfn.envs.preprocessors',
 'gfn.losses',
 'gfn.samplers']

package_data = \
{'': ['*']}

install_requires = \
['einops>=0.6.1',
 'gymnasium>=0.28.1',
 'numpy>=1.21.2',
 'torch>=1.9.0',
 'torchtyping>=0.1.4']

extras_require = \
{'dev': ['pre-commit',
         'pytest',
         'renku-sphinx-theme',
         'sphinx_rtd_theme',
         'sphinx-autoapi',
         'sphinx-math-dollar',
         'sphinx',
         'tox',
         'black==22.3.0',
         'myst-parser'],
 'scripts': ['tqdm', 'wandb', 'simple-parsing==0.0.20']}

setup_kwargs = {
    'name': 'torchgfn',
    'version': '0.1.3',
    'description': 'A torch inplementation of GFlowNets',
    'long_description': '<p align="center">\n    <a>\n\t    <img src=\'https://img.shields.io/badge/python-3.10%2B-blueviolet\' alt=\'Python\' />\n\t</a>\n\t<a href=\'https://gfn.readthedocs.io/en/latest/?badge=latest\'>\n    \t<img src=\'https://readthedocs.org/projects/gfn/badge/?version=latest\' alt=\'Documentation Status\' />\n\t</a>\n    <a>\n\t    <img src=\'https://img.shields.io/badge/code%20style-black-black\' />\n\t</a>\n</p>\n\n</p>\n<p align="center">\n  <a href="https://gfn.readthedocs.io/en/latest/">Documentation</a> ~ <a href="https://github.com/saleml/gfn">Code</a>\n</p>\n\n# gfn: a Python package for GFlowNets\n\n\n\n## Installing the packages\n\nThe codebase requires python >= 3.10\n\n```bash\ngit clone https://github.com/saleml/gfn.git\nconda create -n gfn python=3.10\nconda activate gfn\ncd gfn\npip install .\n```\n\nOptionally, to run scripts, and for [wandb](https://wandb.ai) logging\n\n```bash\npip install .[scripts]\nwandb login\n```\n\n## About this repo\n\nThis repo serves the purpose of fast prototyping [GFlowNet](https://arxiv.org/abs/2111.09266) related algorithms. It decouples the environment definition, the sampling process, and the parametrization used for the GFN loss.\n\nAn example script is provided [here](https://github.com/saleml/gfn/blob/master/scripts/train.py). To run the code, use one of the following:\n\n```bash\npython train.py --env HyperGrid --env.ndim 4 --env.height 8 --n_iterations 100000 --loss TB\npython train.py --env DiscreteEBM --env.ndim 4 --env.alpha 0.5 --n_iterations 10000 --batch_size 64 --temperature 2.\npython train.py --env HyperGrid --env.ndim 2 --env.height 64 --n_iterations 100000 --loss DB --replay_buffer_size 1000 --logit_PB.module_name Uniform --optim sgd --optim.lr 5e-3\npython train.py --env HyperGrid --env.ndim 4 --env.height 8 --env.R0 0.01 --loss FM --optim adam --optim.lr 1e-4\n```\n\n### Example, in a few lines\n\n(⬇️ This example require the [`tqdm`](https://github.com/tqdm/tqdm) package to run. `pip install tqdm` or install all extra requirements with `pip install .[scripts]`)\n\n```python\nimport torch\nfrom tqdm import tqdm\n\nfrom gfn import LogitPBEstimator, LogitPFEstimator, LogZEstimator\nfrom gfn.envs import HyperGrid\nfrom gfn.losses import TBParametrization, TrajectoryBalance\nfrom gfn.samplers import DiscreteActionsSampler, TrajectoriesSampler\n\nif __name__ == "__main__":\n\n    env = HyperGrid(ndim=4, height=8, R0=0.01)  # Grid of size 8x8x8x8\n\n    logit_PF = LogitPFEstimator(env=env, module_name="NeuralNet")\n    logit_PB = LogitPBEstimator(\n        env=env,\n        module_name="NeuralNet",\n        torso=logit_PF.module.torso,  # To share parameters between PF and PB\n    )\n    logZ = LogZEstimator(torch.tensor(0.0))\n\n    parametrization = TBParametrization(logit_PF, logit_PB, logZ)\n\n    actions_sampler = DiscreteActionsSampler(estimator=logit_PF)\n    trajectories_sampler = TrajectoriesSampler(env=env, actions_sampler=actions_sampler)\n\n    loss_fn = TrajectoryBalance(parametrization=parametrization)\n\n    params = [\n        {\n            "params": [\n                val for key, val in parametrization.parameters.items() if "logZ" not in key\n            ],\n            "lr": 0.001,\n        },\n        {"params": [val for key, val in parametrization.parameters.items() if "logZ" in key], "lr": 0.1},\n    ]\n    optimizer = torch.optim.Adam(params)\n\n    for i in (pbar := tqdm(range(1000))):\n        trajectories = trajectories_sampler.sample(n_trajectories=16)\n        optimizer.zero_grad()\n        loss = loss_fn(trajectories)\n        loss.backward()\n        optimizer.step()\n        if i % 25 == 0:\n            pbar.set_postfix({"loss": loss.item()})\n\n```\n\n## Contributing\n\nBefore the first commit:\n\n```bash\npip install .[dev]\npre-commit install\npre-commit run --all-files\n```\n\nRun `pre-commit` after staging, and before committing. Make sure all the tests pass (By running `pytest`).\nThe codebase uses `black` formatter.\n\nTo make the docs locally:\n```bash\ncd docs\nmake html\nopen build/html/index.html\n```\n\n## Details about the codebase\n\n### Defining an environment\n\nA pointed DAG environment (or GFN environment, or environment for short) is a representation for the pointed DAG. The abstract class [Env](https://github.com/saleml/gfn/blob/master/src/gfn/envs/env.py) specifies the requirements for a valid environment definition. To obtain such a representation, the environment needs to specify the following attributes, properties, or methods:\n\n- The `action_space`. Which should be a `gymnasium.spaces.Discrete` object for discrete environments. The last action should correspond to the exit action.\n- The initial state `s_0`, as a `torch.Tensor` of arbitrary dimension.\n- (Optional) The sink state `s_f`, as a `torch.Tensor` of the same shape as `s_0`, used to represent complete trajectories only (within a batch of trajectories of different lengths), and never processed by any model. If not specified, it is set to `torch.full_like(s_0, -float(\'inf\'))`.\n- The method `make_States_class` that creates a subclass of [States](https://github.com/saleml/gfn/blob/master/src/gfn/containers/states.py). The instances of the resulting class should represent a batch of states of arbitrary shape, which is useful to define a trajectory, or a batch of trajectories. `s_0` and `s_f`, along with a tuple called `state_shape` should be defined as class variables, and the subclass (of `States`) should implement masking methods, that specify which actions are possible, in a discrete environment.\n- The methods `maskless_step` and `maskless_backward_step` that specify how an action changes a state (going forward and backward). These functions do not need to handle masking, checking whether actions are allowed, checking whether a state is the sink state, etc... These checks are handled in `Env.step` and `Env.backward_step`\n- The `log_reward` function that assigns a nonnegative reward to every terminating state (i.e. state with all $s_f$ as a child in the DAG). If `log_reward` is not implemented, `reward` needs to be.\n\nIf the states (as represented in the `States` class) need to be transformed to another format before being processed (by neural networks for example), then the environment should define a `preprocessor` attribute, which should be an instance of the [base preprocessor class](https://github.com/saleml/gfn/blob/master/src/gfn/envs/preprocessors/base.py). If no preprocessor is defined, the states are used as is (actually transformed using  [`IdentityPreprocessor`](https://github.com/saleml/gfn/blob/master/src/gfn/envs/preprocessors/base.py), which transforms the state tensors to `FloatTensor`s). Implementing your own preprocessor requires defining the `preprocess` function, and the `output_shape` attribute, which is a tuple representing the shape of *one* preprocessed state.\n\nOptionally, you can define a static `get_states_indices` method that assigns a unique integer number to each state if the environment allows it, and a `n_states` property that returns an integer representing the number of states (excluding $s_f$) in the environment. `get_terminating_states_indices` can also be implemented and serves the purpose of uniquely identifying terminating states of the environment.\n\nFor more details, take a look at [HyperGrid](https://github.com/saleml/gfn/blob/master/src/gfn/envs/hypergrid.py), an environment where all states are terminating states, or at [DiscreteEBM](https://github.com/saleml/gfn/blob/master/src/gfn/envs/discrete_ebm.py), where all trajectories are of the same length but only some states are terminating.\n\n### Other containers\n\nBesides the `States` class, other containers of states are available:\n\n- [Transitions](https://github.com/saleml/gfn/blob/master/src/gfn/containers/transitions.py), representing a batch of transitions $s \\rightarrow s\'$.\n- [Trajectories](https://github.com/saleml/gfn/blob/master/src/gfn/containers/trajectories.py), representing a batch of complete trajectories $\\tau = s_0 \\rightarrow s_1 \\rightarrow \\dots \\rightarrow s_n \\rightarrow s_f$.\n\nThese containers can either be instantiated using a `States` object, or can be initialized as empty containers that can be populated on the fly, allowing the usage of [ReplayBuffer](https://github.com/saleml/gfn/blob/master/src/gfn/containers/replay_buffer.py)s.\n\nThey inherit from the base `Container` [class](https://github.com/saleml/gfn/blob/master/src/gfn/containers/base.py), indicating some helpful methods.\n\nIn most cases, one needs to sample complete trajectories. From a batch of trajectories, a batch of states and batch of transitions can be defined using `Trajectories.to_transitions()` and `Trajectories.to_states()`. These exclude meaningless transitions and states that were added to the batch of trajectories to allow for efficient batching.\n\n### Estimators and Modules\n\nTraining GFlowNets requires one or multiple estimators. As of now, only discrete environments are handled. All estimators are subclasses of [FunctionEstimator](https://github.com/saleml/gfn/blob/master/src/gfn/estimators.py), implementing a `__call__` function that takes as input a batch of [States](https://github.com/saleml/gfn/blob/master/src/gfn/containers/states.py).\n\n- [LogEdgeFlowEstimator](https://github.com/saleml/gfn/blob/master/src/gfn/estimators.py). It outputs a `(*batch_shape, n_actions)` tensor representing $\\log F(s \\rightarrow s\')$, including when $s\' = s_f$.\n- [LogStateFlowEstimator](https://github.com/saleml/gfn/blob/master/src/gfn/estimators.py). It outputs a `(*batch_shape, 1)` tensor representing $\\log F(s)$. When used with `forward_looking=True`, $\\log F(s)$ is parametrized as the sum of a function approximator and $\\log R(s)$ - which is only possible for environments where all states are terminating.\n- [LogitPFEstimator](https://github.com/saleml/gfn/blob/master/src/gfn/estimators.py). It outputs a `(*batch_shape, n_actions)` tensor representing $logit(s\' \\mid s)$, such that $P_F(s\' \\mid s) = softmax_{s\'}\\ logit(s\' \\mid s)$, including when $s\' = s_f$.\n- [LogitPBEstimator](https://github.com/saleml/gfn/blob/master/src/gfn/estimators.py). It outputs a `(*batch_shape, n_actions - 1)` tensor representing $logit(s\' \\mid s)$, such that $P_B(s\' \\mid s) = softmax_{s\'}\\ logit(s\' \\mid s)$.\n\nDefining an estimator requires the environment, and a [module](https://github.com/saleml/gfn/blob/master/src/gfn/modules.py) instance. Modules inherit from the [GFNModule](https://github.com/saleml/gfn/blob/master/src/gfn/modules.py) class, which can be seen as an extension of `torch.nn.Module`. Alternatively, a module is created by providing which module type to use (e.g. "NeuralNet" or "Uniform" or "Zero"). A Basic MLP is provided as the [NeuralNet](https://github.com/saleml/gfn/blob/master/src/gfn/modules.py) class, but any function approximator should be possible.\n\nSaid differently, a `States` object is first transformed via the environment\'s preprocessor to a `(*batch_shape, *output_shape)` float tensor. The preprocessor\'s output shape should match the module input shape (if any). The preprocessed states are then passed as inputs to the module, returning the desired output (either flows or probabilities over children in the DAG).\n\nEach module has a `named_parameters` functions that returns a dictionary of the learnable parameters. This attribute is transferred to the corresponding estimator.\n\nAdditionally, a [LogZEstimator](https://github.com/saleml/gfn/blob/master/src/gfn/estimators.py) is provided, which is a container for a scalar tensor representing $\\log Z$, the log-partition function, useful for the Trajectory Balance loss for example. This estimator also has a `named_parameters` function.\n\n### Samplers\n\nAn [ActionsSampler](https://github.com/saleml/gfn/blob/master/src/gfn/samplers/actions_samplers.py) object defines how actions are sampled at each state of the DAG. As of now, only [DiscreteActionsSampler](https://github.com/saleml/gfn/blob/master/src/gfn/samplers/actions_samplers.py)s are implemented. The require an estimator (of $P_F$, $P_B$, or edge flows) defining the action probabilities. These estimators can contain any type of modules (including random action sampling for example). A [BackwardDiscreteActionsSampler](https://github.com/saleml/gfn/blob/master/src/gfn/samplers/actions_samplers.py) class is provided to sample parents of a state, which is helpful to sample trajectories starting from their last states.\n\nThey are at the core of [TrajectoriesSampler](https://github.com/saleml/gfn/blob/master/src/gfn/samplers/trajectories_sampler.py)s, which implements the `sample_trajectories` method, that sample a batch of trajectories starting from a given set of initial states or starting from $s_0$.\n\n### Losses\n\nGFlowNets can be trained with different losses, each of which requires a different parametrization. A parametrization is a dataclass, which can be seen as a container of different estimators. Each parametrization defines a distribution over trajectories, via the `parametrization.Pi` method, and a distribution over terminating states, via the `parametrization.P_T` method. Both distributions should be instances of the classes defined [here](https://github.com/saleml/gfn/blob/master/src/gfn/distributions.py).\n\nThe base classes for losses and parametrizations are provided [here](https://github.com/saleml/gfn/blob/master/src/gfn/losses/base.py).\n\nCurrently, the implemented losses are:\n\n- Flow Matching\n- Detailed Balance\n- Trajectory Balance\n- Sub-Trajectory Balance. By default, each sub-trajectory is weighted geometrically (within the trajectory) depending on its length. This corresponds to the strategy defined [here](https://www.semanticscholar.org/reader/f2c32fe3f7f3e2e9d36d833e32ec55fc93f900f5). Other strategies exist and are implemented [here](https://github.com/saleml/gfn/blob/master/src/gfn/losses/sub_trajectory_balance.py).\n- Log Partition Variance loss. Introduced [here](https://arxiv.org/abs/2302.05446)\n\n### Solving for the flows using Dynamic Programming\n\nA simple script that propagates trajectories rewards through the DAG to define edge flows in a deterministic way (by visiting each edge once only) is provided [here](https://github.com/saleml/gfn/blob/master/scripts/dynamic_programming.py). Do not use the script on large environments !\n',
    'author': 'Salem Lahou',
    'author_email': 'salemlahlou9@gmail.com',
    'maintainer': 'None',
    'maintainer_email': 'None',
    'url': 'None',
    'package_dir': package_dir,
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'extras_require': extras_require,
    'python_requires': '>=3.10,<4.0',
}


setup(**setup_kwargs)
