tf_agents.utils.numpy_storage.NumpyStorage

View source on GitHub

A class to store nested objects in a collection of numpy arrays.

If a data_spec of {'foo': ArraySpec(shape=(4,), dtype=np.uint8), 'bar': ArraySpec(shape=(3, 7), dtype=np.float32)} were used, then this would create two arrays, one for the 'foo' key and one for the 'bar' key. The .get and .set methods would return/take Python dictionaries, but break down the component arrays before storing them.

data_spec An ArraySpec or a list/tuple/nest of ArraySpecs describing a single item that can be stored in this table.
capacity The maximum number of items that can be stored in the buffer.

ValueError If data_spec is not an instance or nest of ArraySpecs.

name Returns the name of this module as passed or determined in the ctor.

name_scope Returns a tf.name_scope instance for this class.
submodules Sequence of all sub-modules.

Submodules are modules which are properties of this module, or found as properties of modules which are properties of this module (and so on).

a = tf.Module()
b = tf.Module()
c = tf.Module()
a.b = b
b.c = c
list(a.submodules) == [b, c]
True
list(b.submodules) == [c]
True
list(c.submodules) == []
True

trainable_variables Sequence of trainable variables owned by this module and its submodules.

variables Sequence of variables owned by this module and its submodules.

Methods

get

View source

Get value stored at idx.

set

View source

Set table_idx to value.

with_name_scope

Decorator to automatically enter the module name scope.

class MyModule(tf.Module):
  @tf.Module.with_name_scope
  def __call__(self, x):
    if not hasattr(self, 'w'):
      self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
    return tf.matmul(x, self.w)

Using the above module would produce tf.Variables and tf.Tensors whose names included the module name:

mod = MyModule()
mod(tf.ones([1, 2]))
<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
mod.w
<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
numpy=..., dtype=float32)>

Args
method The method to wrap.

Returns
The original method wrapped such that it enters the module's name scope.