Tensor#

The genesis.grad.Tensor class extends torch.Tensor to support end-to-end gradient flow through Genesis World simulations.

Overview#

Genesis World Tensors:

  • Extend PyTorch tensors with scene tracking

  • Enable automatic gradient propagation through physics

  • Support all standard PyTorch operations

  • Track parent tensors for gradient flow

Usage#

Genesis World Tensors are automatically created when you access state:

import genesis as gs
import torch

gs.init()

scene = gs.Scene(
    sim_options=gs.options.SimOptions(
        requires_grad=True,
    ),
)

robot = scene.add_entity(gs.morphs.URDF(file="robot.urdf"))
scene.build()

# These return genesis.grad.Tensor
pos = robot.get_pos()       # Genesis World Tensor
vel = robot.get_vel()       # Genesis World Tensor
qpos = robot.get_qpos()     # Genesis World Tensor

Gradient flow#

# Forward pass
scene.step()
pos = robot.get_pos()

# Compute loss
target = torch.tensor([1.0, 0.0, 0.5], device=gs.device)
loss = (pos - target).pow(2).sum()

# Backward pass - flows through simulation
loss.backward()

Detaching from scene#

To prevent gradients from flowing through the simulation:

# Detach and remove scene tracking
pos_detached = pos.detach(sceneless=True)

# Or explicitly
pos_sceneless = pos.sceneless()

Checking scene association#

# Check if tensor is associated with a scene
if pos.scene is not None:
    print(f"Tensor from scene: {pos.scene.uid}")

API reference#

class genesis.grad.tensor.Tensor(*args, scene=None, **kwargs)[source]#

Bases: Tensor

This is the genesis customization of torch’s Tensor datatype. This allows our customizations, a few safety checks and also more elegant end-to-end gradient flow.

detach(*args, sceneless=True, **kwargs)[source]#

Returns a new Tensor, detached from the current graph.

The result will never require gradient.

This method also affects forward mode AD gradients and the result will never have forward mode AD gradients.

Note

Returned Tensor shares the same storage with the original one. In-place modifications on either of them will be seen, and may trigger errors in correctness checks.

backward(*args, **kwargs)[source]#

Computes the gradient of current tensor wrt graph leaves.

The graph is differentiated using the chain rule. If the tensor is non-scalar (i.e. its data has more than one element) and requires gradient, the function additionally requires specifying a gradient. It should be a tensor of matching type and shape, that represents the gradient of the differentiated function w.r.t. self.

This function accumulates gradients in the leaves - you might need to zero .grad attributes or set them to None before calling it. See Default gradient layouts for details on the memory layout of accumulated gradients.

Note

If you run any forward ops, create gradient, and/or call backward in a user-specified CUDA stream context, see Stream semantics of backward passes.

Note

When inputs are provided and a given input is not a leaf, the current implementation will call its grad_fn (though it is not strictly needed to get this gradients). It is an implementation detail on which the user should not rely. See pytorch/pytorch#60521 for more details.

Parameters:

  • gradient (Tensor, optional) – The gradient of the function being differentiated w.r.t. self. This argument can be omitted if self is a scalar. Defaults to None.

  • retain_graph (bool, optional) – If False, the graph used to compute the grads will be freed; If True, it will be retained. The default is None, in which case the value is inferred from create_graph (i.e., the graph is retained only when higher-order derivative tracking is requested). Note that in nearly all cases setting this option to True is not needed and often can be worked around in a much more efficient way.

  • create_graph (bool, optional) – If True, graph of the derivative will be constructed, allowing to compute higher order derivative products. Defaults to False.

  • inputs (Sequence[Tensor] or dict[str, Tensor], optional) – Inputs w.r.t. which the gradient will be accumulated into .grad. All other tensors will be ignored. If not provided, the gradient is accumulated into all the leaf Tensors that were used to compute the tensors. A dict of tensors (e.g. dict(model.named_parameters())) is also accepted. Defaults to None.

zero_grad()[source]#

A handy method that resembles nn.Module.zero_grad().

sceneless()[source]#

Return a tensor detached from a scene.

assert_contiguous()[source]#
assert_sceneless()[source]#

See also#