gfn.actions
Module Contents
Classes
Base class for actions for all GFlowNet environments. |
- class gfn.actions.Actions(tensor)
Bases:
abc.ABCBase class for actions for all GFlowNet environments.
Each environment needs to subclass this class. A generic subclass for discrete actions with integer indices is provided Note that all actions need to have the same shape.
- Parameters
tensor (torchtyping.TensorType[batch_shape, action_shape]) –
- batch_shape
the batch_shape from the input tensor.
- property device: torch.device
- Return type
torch.device
- property is_dummy: torchtyping.TensorType[batch_shape, torch.bool]
Returns a boolean tensor indicating whether the actions are dummy actions.
- Return type
torchtyping.TensorType[batch_shape, torch.bool]
- property is_exit: torchtyping.TensorType[batch_shape, torch.bool]
Returns a boolean tensor indicating whether the actions are exit actions.
- Return type
torchtyping.TensorType[batch_shape, torch.bool]
- action_shape :ClassVar[tuple[int, Ellipsis]]
- dummy_action :ClassVar[torchtyping.TensorType[Actions.action_shape]]
- exit_action :ClassVar[torchtyping.TensorType[Actions.action_shape]]
- __len__()
- Return type
int
- __repr__()
Return repr(self).
- __setitem__(index, actions)
Set particular actions of the batch.
- Parameters
index (int | Sequence[int] | Sequence[bool]) –
actions (Actions) –
- Return type
None
- compare(other)
Compares the actions to a tensor of actions.
- Parameters
other (torchtyping.TensorType[batch_shape, action_shape]) – tensor of actions
- Return type
torchtyping.TensorType[batch_shape, torch.bool]
- Returns: boolean tensor of shape batch_shape indicating whether the actions are
equal.
- extend(other)
Collates to another Actions object of the same batch shape.
- Parameters
other (Actions) –
- Return type
None
- extend_with_dummy_actions(required_first_dim)
Extends an Actions instance along the first dimension with dummy actions.
The Actions instance batch_shape must be 2-dimensional. This is used to pad trajectories actions.
- Parameters
required_first_dim (int) – the target size of the first dimension post expansion.
- Return type
None
- classmethod make_dummy_actions(batch_shape)
Creates an Actions object of dummy actions with the given batch shape.
- Parameters
batch_shape (tuple[int]) –
- Return type
- classmethod make_exit_actions(batch_shape)
Creates an Actions object of exit actions with the given batch shape.
- Parameters
batch_shape (tuple[int]) –
- Return type
- classmethod stack(actions_list)
Stacks a list of Actions objects into a single Actions object.
The individual actions need to have the same batch shape. An example application is when the individual actions represent per-step actions of a batch of trajectories (in which case, the common batch_shape would be (n_trajectories,), and the resulting Actions object would have batch_shape (n_steps, n_trajectories).