reference/generated/torchrl.render.extract_qpos_trajectory

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

extract_qpos_trajectory

torchrl.render.extract_qpos_trajectory(rollout: TensorDictBase, qpos_key: NestedKey | str = 'qpos') → list[list[float]]

Extracts a qpos trajectory from a rollout TensorDict.

Parameters:

• rollout – Rollout TensorDict saved by rlrender.

• qpos_key – TensorDict key containing a T x nq qpos tensor.

Returns:

A Python list of qpos waypoints suitable for play_mujoco_wasm_trajectory().

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> from torchrl.render.mujoco_wasm import extract_qpos_trajectory
>>> rollout = TensorDict({"qpos": torch.zeros(2, 3)}, batch_size=[2])
>>> extract_qpos_trajectory(rollout, "qpos")
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]