xdrl 🔍

Use TDHook methods on native TorchRL modules and losses.

Getting Started

pip install xdrl

import torch
from tensordict import TensorDict
from tensordict.nn import TensorDictModule
from tdhook.latent import ActivationCaching
from tdhook.workflow import Workflow
from xdrl import interpret

model = TensorDictModule(
    torch.nn.Sequential(
        torch.nn.Linear(4, 8),
        torch.nn.Tanh(),
        torch.nn.Linear(8, 2),
    ),
    in_keys=["observation"],
    out_keys=["action"],
)
policy = interpret(model)
batch = TensorDict({"observation": torch.randn(8, 4)}, batch_size=[8])
result = policy.run(
    Workflow(ActivationCaching("module.1", cache_key=("activations", "hidden"))),
    batch,
)

assert result.data["action"].shape == (8, 2)
assert result.data["activations", "hidden", "module.1"].shape == (8, 8)

interpret preserves the TensorDict API and adds .run(...) for TDHook workflows. The policy, TensorDict keys, and execution behavior stay native to TorchRL; XDRL only supplies the interpretability view.

For recurrent TorchRL modules, see RecurrentSemantics in the API reference.

Development

This project uses uv to manage Python dependencies and just to run the test and documentation gates.

Documentation

• Getting Started
• Tutorials
• API Reference
• About

License

xdrl is licensed under the MIT License. See LICENSE for details.