TensorDictBase.backward#

TensorDictBase.backward(gradient: TensorDictBase | None = None, retain_graph: bool | None = None, create_graph: bool = False, inputs: TensorDictBase | Sequence[Tensor] | None = None) → None#

Computes the gradients of the leaves of the tensordict that require gradients.

Like torch.Tensor.backward() for every leaf at once: the leaves that require gradients, nested ones included, are differentiated in a single torch.autograd.backward() call, and the others are ignored. The gradients are accumulated in the .grad of the graph leaves, as usual. Returns None.

Without gradient, every leaf that requires gradients must have a single element, and its gradient is 1: td.backward() is then the same as td.sum(reduce=True).backward().

Parameters:

gradient (TensorDictBase, optional) – the gradient of each leaf, under the same nested key and with the same shape. The leaves can differ in shape, dtype and device. Defaults to None, which needs single-element leaves.
retain_graph (bool, optional) – if False, the graph used to compute the gradients is freed. Defaults to None, which takes the value of create_graph.
create_graph (bool, optional) – if True, the graph of the derivative is built, for higher-order derivatives. Defaults to False.
inputs (TensorDictBase or sequence of Tensor, optional) – the tensors whose .grad receives the gradients; the others are ignored. Defaults to None, which accumulates into all the graph leaves used to compute the leaves of the tensordict.

Raises:

RuntimeError – if no leaf requires gradients, if gradient is not given and a leaf has more than one element, or if a gradient has the wrong shape.
KeyError – if gradient has no entry for a leaf that requires gradients.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> x = torch.ones(3, requires_grad=True)
>>> losses = TensorDict(actor_loss=x.sum(), critic_loss=x.pow(2).sum())
>>> losses.backward()  # a gradient of 1 for each single-element leaf
>>> x.grad
tensor([3., 3., 3.])

Weight the losses with a gradient tensordict:

>>> x.grad = None
>>> losses = TensorDict(actor_loss=x.sum(), critic_loss=x.pow(2).sum())
>>> weights = TensorDict(actor_loss=torch.tensor(0.5), critic_loss=torch.tensor(1.0))
>>> losses.backward(weights)
>>> x.grad
tensor([2.5000, 2.5000, 2.5000])

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.

TensorDictBase.backward#

Docs

Tutorials

Resources