reference/render

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

Rendering applications

torchrl.render provides reusable utilities behind the rlrender and torchrl-render commands. The command-line entry point imports trusted user policy and environment factories, loads a local checkpoint, collects one or more rollouts, captures RGB frames from TensorDict pixels or env.render(), and writes a reproducible artifact.

Notebook artifacts can also include an optional MuJoCo WASM sidecar viewer. In that mode, the notebook imports helper functions from torchrl.render to start a local Vite viewer, load an MJCF scene in browser-side MuJoCo, and stream qpos trajectories into the live iframe. When the environment exposes native MuJoCo state or wraps a Gymnasium MuJoCo environment, MujocoStateReader records qpos directly from the simulator rather than deriving it from policy observations. By default, rlrender collects and saves trajectories before writing the notebook. Use --notebook-rollout-mode live to write a notebook that constructs the policy and environment inside the kernel and generates trajectories when its cells are run. Use --notebook-rollout-mode both to save collected rollouts and also keep an in-notebook collection cell. The generated notebook should stay thin: reusable display, playback, rollout, and acknowledgement helpers live in TorchRL rather than being copied into each notebook.

The MuJoCo WASM viewer requires Node.js and either npm or pnpm. The viewer installs the generated Vite project’s pinned JavaScript dependencies when node_modules is absent, which requires network access. The generated node_modules directory is reused when present.

Factories can be addressed as module.submodule:callable or as a local file path such as /path/to/render_factories.py:make_env. The base TorchRL package does not install video or image encoders for this feature. Use the optional rendering dependencies when writing MP4, GIF, PNG, or YAML-backed configs:

uv run --extra rendering rlrender --help

Core API

render_policy(config)	Renders a policy according to config and writes the requested artifact.
make_render_env(config[, checkpoint])	Builds and prepares an environment for rendering.
load_render_policy(config[, env, ...])	Builds and prepares a policy for rendering.
collect_render_rollouts(env, policy, config)	Collects sequential render rollouts.
write_render_artifact(result, config)	Writes the configured render artifact and sidecar metadata.
import_from_string(spec)	Imports an object from a "module:attribute" string.
call_with_supported_kwargs(factory, ...)	Calls a user factory with a spec object or supported keyword arguments.
load_checkpoint(path[, map_location, ...])	Loads a local PyTorch checkpoint.
save_render_checkpoint(path, model, *[, ...])	Writes a checkpoint in the layout expected by rlrender factories.
checkpoint_hash(path)	Compute a SHA256 digest over all checkpoint bytes.
infer_state_dict(payload[, key])	Infers a model state dict from common checkpoint payload layouts.
parse_nested_key(value)	Parses dotted strings into TensorDict nested keys.
key_to_string(key)	Formats a TensorDict nested key for config and metadata output.
normalize_policy(policy, config)	Normalizes a policy into a TensorDict-compatible callable.
add_step_counter(env, max_steps)	Adds a StepCounter when supported.
seed_env(env, seed)	Seeds an environment if it exposes a known seed method.
normalize_env(env, config)	Normalizes external environments into TorchRL wrappers when feasible.
write_mujoco_wasm_viewer(output_dir, ...[, ...])	Writes a local Vite viewer for MuJoCo WASM notebook rendering.
display_mujoco_wasm_viewer(viewer_dir, *[, ...])	Starts and displays a generated MuJoCo WASM viewer in a notebook.
send_mujoco_wasm_qpos(qpos, *[, port, ...])	Sends one qpos vector to a live MuJoCo WASM notebook viewer.
play_mujoco_wasm_trajectory(qpos, *[, dt, ...])	Plays a qpos trajectory in a live MuJoCo WASM notebook viewer.
extract_qpos_trajectory(rollout[, qpos_key])	Extracts a qpos trajectory from a rollout TensorDict.

Configuration and results

RenderConfig(ckpt, policy, ~typing.Any], ...)	Configuration for rendering policy rollouts.
RenderEnvSpec(device, seed, max_steps, ...)	Context object passed to environment factories.
RenderPolicySpec(ckpt_path, checkpoint, ...)	Context object passed to policy factories.
RenderResult(artifact_path, trajectories, ...)	Result returned by torchrl.render.render_policy().
FrameBundle(frames, ~typing.Any], step, ...)	One rendered step containing one or more named camera frames.
TensorDictPolicyAdapter(policy, obs_key, ...)	Adapts plain tensor policies to a TensorDict policy callable.

Backends

RenderBackend(*args, **kwargs)	Protocol implemented by rlrender frame-capture backends.
MujocoStateReader()	Reads simulator state from TorchRL-native and Gym MuJoCo environments.
TensorDictPixelsBackend()	Captures frames from TensorDict pixel entries.
EnvRenderBackend()	Captures frames by calling env.render().
NullRenderBackend()	Fallback backend used when no RGB renderer is available.

Lower-level helpers

build_parser()	Builds the rlrender command-line parser.
config_from_args(args)	Constructs a RenderConfig from parsed CLI args.
main([argv])	Entry point for rlrender and torchrl-render.

build_notebook(result, config)	Builds a minimal reproducible render report notebook.
write_render_notebook(result, config, path)	Writes a render report notebook as plain ipynb JSON.

normalize_frame(frame)	Converts a tensor-like image into an H x W x 3 uint8 array.
normalize_frame_output(output)	Normalizes renderer output into named RGB frames.
compose_frame_grid(frames[, layout])	Composes multiple frames into one RGB image.
encode_video(frames, path, fps, *[, video_codec])	Encodes RGB frames as an MP4 using TorchRL's torchcodec writer.
encode_gif(frames, path, fps)	Encodes RGB frames as an animated GIF using Pillow.
write_png(frame, path)	Writes one RGB frame as a PNG file using Pillow.