Skip to content

Latest commit

History

History
75 lines (61 loc) 路 2.75 KB

File metadata and controls

75 lines (61 loc) 路 2.75 KB
logo

xdrl 馃攳

Documentation xdrl license uv Ruff python versions

codecov ci docs

Typed, inspectable TorchRL model interactions with TDHook observability and intervention.

Getting Started

pip install xdrl
import torch
from tensordict import TensorDict
from tensordict.nn import TensorDictModule
from xdrl import (
    BatchSemantics,
    KeyPresence,
    KeyRole,
    KeySchema,
    ModelRole,
    TensorDictSchema,
    validate_module,
)

batch = TensorDict({"observation": torch.randn(8, 4)}, batch_size=[8])
policy = TensorDictModule(
    torch.nn.Linear(4, 2),
    in_keys=["observation"],
    out_keys=["action"],
)
policy.role = ModelRole.ACTOR
batch_dims = BatchSemantics(("env",))
policy.input_schema = TensorDictSchema(
    (KeySchema("observation", KeyRole.OBSERVATION, KeyPresence.REQUIRED),),
    batch_dims,
)
policy.output_schema = TensorDictSchema(
    (KeySchema("action", KeyRole.ACTION, KeyPresence.PRODUCED),),
    batch_dims,
)

result = validate_module(policy, batch)
assert result["action"].shape == (8, 2)

Development

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

Documentation

License

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