tf_agents.utils.nest_utils.batch_nested_tensors

Add batch dimension if needed to nested tensors while checking their specs.

tf_agents.utils.nest_utils.batch_nested_tensors(
    tensors, specs=None
)

If specs is None, a batch dimension is added to each tensor. If specs are provided, each tensor is compared to the corresponding spec, and a batch dimension is added only if the tensor doesn't already have it.

For each tensor, it checks the dimensions with respect to specs, and adds an extra batch dimension if it doesn't already have it.

tensors Nested list/tuple or dict of Tensors.
specs Nested list/tuple or dict of TensorSpecs, describing the shape of the non-batched Tensors.

A nested batched version of each tensor.

ValueError if the tensors and specs have incompatible dimensions or shapes.