`popgym.baselines.ray_models.ray_diffnc`

Module Contents

Classes

DiffNC

The differentiable neural computer from

class popgym.baselines.ray_models.ray_diffnc.DiffNC(obs_space: gymnasium.spaces.Space, action_space: gymnasium.spaces.Space, num_outputs: int, model_config: ray.rllib.utils.typing.ModelConfigDict, name: str, **custom_model_kwargs)

Bases: popgym.baselines.ray_models.base_model.BaseModel

The differentiable neural computer from

@techreport{wayne_unsupervised_2018,
    title = {Unsupervised {Predictive} {Memory} in a {Goal}-{Directed} {Agent}},
    url = {http://arxiv.org/abs/1803.10760},
    number = {arXiv:1803.10760},
    urldate = {2022-09-09},
    institution = {arXiv},
    month = mar,
    year = {2018},
    doi = {10.48550/arXiv.1803.10760},
}

MODEL_CONFIG

MEMORY_KEYS = ['memory', 'link_matrix', 'precedence', 'read_weights', 'write_weights', 'usage_vector']

initial_state() → List[ray.rllib.utils.typing.TensorType]

Return the initial states for your memory model.

The shape of states returned here should NOT contain the batch dimension. The batch dimension will be prepended by RLlib, depending on the number of episodes/rollouts in the batch

Returns:: List of tensors denoting the t=0 recurrent state

unpack_state(state: List[ray.rllib.utils.typing.TensorType]) → Tuple[List[Tuple[ray.rllib.utils.typing.TensorType, ray.rllib.utils.typing.TensorType]], Dict[str, ray.rllib.utils.typing.TensorType], ray.rllib.utils.typing.TensorType]: Given a list of tensors, reformat for self.dnc input

pack_state(ctrl_hidden: List[Tuple[ray.rllib.utils.typing.TensorType, ray.rllib.utils.typing.TensorType]], memory_dict: Dict[str, ray.rllib.utils.typing.TensorType], read_vecs: ray.rllib.utils.typing.TensorType) → List[ray.rllib.utils.typing.TensorType]: Given the dnc output, pack it into a list of tensors for rllib state. Order is ctrl_hidden, read_vecs, memory_dict

build_dnc(device_idx: int | None) → torch.nn.Module

forward_memory(z: ray.rllib.utils.typing.TensorType, state: List[ray.rllib.utils.typing.TensorType], t_starts: ray.rllib.utils.typing.TensorType, seq_lens: ray.rllib.utils.typing.TensorType) → Tuple[ray.rllib.utils.typing.TensorType, List[ray.rllib.utils.typing.TensorType]]

Forward for your custom memory model.

Args:

z: Preprocessed features of shape [B, T, F], with padding along the time: dimension

state: Recurrent states of shape [B, …] t_starts: Tensor of size [B] denoteint the length of the rollout so far.

Note that in some cases RLlib might chunk a long rollout into multiple forward passes. This tracks the length throughout all forward passes.

seq_lens: Tensor of size [B] denoting the number of non-padding elements: in z. E.g. seq_lens == [100, 60], z.shape[1] == 128 means the first 100 elements of the first batch dimension are valid and the first 60 elements of the second batch dimension are valid. Rest are padding.

Returns:

(output, state): where output is [B, T, D] and state is [B, …]. Note that the padding must be present in the output. The state must be exactly the same shape as the input state.

popgym.baselines.ray_models.ray_diffnc

Module Contents

Classes

`popgym.baselines.ray_models.ray_diffnc`