Skip to content

core.base_abstractions.misc#

[view_source]

Memory#

class Memory(Dict)

[view_source]

Memory.check_append#

 | check_append(key: str, tensor: torch.Tensor, sampler_dim: int) -> "Memory"

[view_source]

Appends a new memory type given its identifier, its memory tensor and its sampler dim.

Parameters

  • key: string identifier of the memory type
  • tensor: memory tensor
  • sampler_dim: sampler dimension

Returns

Updated Memory

Memory.tensor#

 | tensor(key: str) -> torch.Tensor

[view_source]

Returns the memory tensor for a given memory type.

Parameters

  • key: string identifier of the memory type

Returns

Memory tensor for type key

Memory.sampler_dim#

 | sampler_dim(key: str) -> int

[view_source]

Returns the sampler dimension for the given memory type.

Parameters

  • key: string identifier of the memory type

Returns

The sampler dim

Memory.sampler_select#

 | sampler_select(keep: Sequence[int]) -> "Memory"

[view_source]

Equivalent to PyTorch index_select along the sampler_dim of each memory type.

Parameters

  • keep: a list of sampler indices to keep

Returns

Selected memory

Memory.set_tensor#

 | set_tensor(key: str, tensor: torch.Tensor) -> "Memory"

[view_source]

Replaces tensor for given key with an updated version.

Parameters

  • key: memory type identifier to update
  • tensor: updated tensor

Returns

Updated memory

Memory.step_select#

 | step_select(step: int) -> "Memory"

[view_source]

Equivalent to slicing with length 1 for the step (i.e first) dimension in rollouts storage.

Parameters

  • step: step to keep

Returns

Sliced memory with a single step

Memory.step_squeeze#

 | step_squeeze(step: int) -> "Memory"

[view_source]

Equivalent to simple indexing for the step (i.e first) dimension in rollouts storage.

Parameters

  • step: step to keep

Returns

Sliced memory with a single step (and squeezed step dimension)

Memory.slice#

 | slice(dim: int, start: Optional[int] = None, stop: Optional[int] = None, step: int = 1) -> "Memory"

[view_source]

Slicing for dimensions that have same extents in all memory types. It also accepts negative indices.

Parameters

  • dim: the dimension to slice
  • start: the index of the first item to keep if given (default 0 if None)
  • stop: the index of the first item to discard if given (default tensor size along dim if None)
  • step: the increment between consecutive indices (default 1)

Returns

Sliced memory