gfn.envs.hypergrid

Copied and Adapted from https://github.com/Tikquuss/GflowNets_Tutorial

Module Contents

Classes

HyperGrid

Base class for environments, showing which methods should be implemented.

Attributes

BackwardMasksTensor

ForwardMasksTensor

OneStateTensor

StatesTensor

TensorBool

TensorFloat

TensorLong

preprocessors_dict

gfn.envs.hypergrid.BackwardMasksTensor
gfn.envs.hypergrid.ForwardMasksTensor
class gfn.envs.hypergrid.HyperGrid(ndim=2, height=4, R0=0.1, R1=0.5, R2=2.0, reward_cos=False, device_str='cpu', preprocessor_name='KHot')

Bases: gfn.envs.env.Env

Base class for environments, showing which methods should be implemented. A common assumption for all environments is that all actions are discrete, represented by a number in {0, …, n_actions - 1}, the last one being the exit action.

Parameters

  • ndim (int) –

  • height (int) –

  • R0 (float) –

  • R1 (float) –

  • R2 (float) –

  • reward_cos (bool) –

  • device_str (Literal[cpu, cuda]) –

  • preprocessor_name (Literal[KHot, OneHot, Identity]) –

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

build_grid()

Utility function to build the complete grid

Return type

gfn.containers.states.States

get_states_indices(states)
Parameters

states (gfn.containers.states.States) –

Return type

TensorLong

get_terminating_states_indices(states)
Parameters

states (gfn.containers.states.States) –

Return type

TensorLong

is_exit_actions(actions)

Returns True if the action is an exit action.

Parameters

actions (TensorLong) –

Return type

TensorBool

log_reward(final_states)

Either this or reward needs to be implemented.

Parameters

final_states (gfn.containers.states.States) –

Return type

TensorFloat

make_States_class()

Creates a States class for this environment

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 (TensorLong) –

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 (TensorLong) –

Return type

None

true_reward(final_states)
Parameters

final_states (gfn.containers.states.States) –

Return type

TensorFloat

gfn.envs.hypergrid.OneStateTensor
gfn.envs.hypergrid.StatesTensor
gfn.envs.hypergrid.TensorBool
gfn.envs.hypergrid.TensorFloat
gfn.envs.hypergrid.TensorLong
gfn.envs.hypergrid.preprocessors_dict