gfn.envs.discrete_ebm

Module Contents

Classes

DiscreteEBMEnv	Environment for discrete energy-based models, based on https://arxiv.org/pdf/2202.01361.pdf
EnergyFunction	Base class for energy functions
IsingModel	Ising model energy function

Attributes

BackwardMasksTensor
BatchTensor
ForwardMasksTensor
IsingJTensor
StatesTensor

gfn.envs.discrete_ebm.BackwardMasksTensor

gfn.envs.discrete_ebm.BatchTensor

class gfn.envs.discrete_ebm.DiscreteEBMEnv(ndim, energy=None, alpha=1.0, device_str='cpu')

Bases: gfn.envs.env.Env

Environment for discrete energy-based models, based on https://arxiv.org/pdf/2202.01361.pdf

Parameters

• ndim (int) –

• energy (EnergyFunction | None) –

• alpha (float) –

• device_str (Literal[cpu, cuda]) –

property all_states: gfn.containers.states.States

Returns a batch of all states for environments with enumerable states. The batch_shape should be (n_states,). This should satisfy: self.get_states_indices(self.all_states) == torch.arange(self.n_states)

Return type

gfn.containers.states.States

property log_partition: float

Returns the logarithm of the partition function.

Return type

float

property n_states: int

Return type

int

property n_terminating_states: int

Return type

int

property terminating_states: gfn.containers.states.States

Returns a batch of all terminating states for environments with enumerable states. The batch_shape should be (n_terminating_states,). This should satisfy: self.get_terminating_states_indices(self.terminating_states) == torch.arange(self.n_terminating_states)

Return type

gfn.containers.states.States

property true_dist_pmf: torch.Tensor

Returns a one-dimensional tensor representing the true distribution.

Return type

torch.Tensor

get_states_indices(states)

The chosen encoding is the following: -1 -> 0, 0 -> 1, 1 -> 2, then we convert to base 3

Parameters

states (gfn.containers.states.States) –

Return type

BatchTensor

get_terminating_states_indices(states)

Parameters

states (gfn.containers.states.States) –

Return type

BatchTensor

is_exit_actions(actions)

Returns True if the action is an exit action.

Parameters

actions (BatchTensor) –

Return type

BatchTensor

log_reward(final_states)

Either this or reward needs to be implemented.

Parameters

final_states (gfn.containers.states.States) –

Return type

BatchTensor

make_States_class()

Returns a class that inherits from States and implements the environment-specific methods.

Return type

type[gfn.containers.states.States]

maskless_backward_step(states, actions)

Same as the backward_step function, but without worrying whether or not the actions are valid, or masking.

Parameters

• states (StatesTensor) –

• actions (BatchTensor) –

Return type

None

maskless_step(states, actions)

Same as the step function, but without worrying whether or not the actions are valid, or masking.

Parameters

• states (StatesTensor) –

• actions (BatchTensor) –

Return type

None

class gfn.envs.discrete_ebm.EnergyFunction

Bases: torch.nn.Module, abc.ABC

Base class for energy functions

abstract forward(states)

Parameters

states (StatesTensor) –

Return type

BatchTensor

gfn.envs.discrete_ebm.ForwardMasksTensor

gfn.envs.discrete_ebm.IsingJTensor

class gfn.envs.discrete_ebm.IsingModel(J)

Bases: EnergyFunction

Ising model energy function

Parameters

J (IsingJTensor) –

forward(states)

Parameters

states (StatesTensor) –

Return type

BatchTensor

gfn.envs.discrete_ebm.StatesTensor