pytorch / pytorch/rl

[Feature Request] Transformers for first-class support in torchRL

Open
#4,192 1 comment 0 reactions 1 assignee View on GitHub

@vmoens is already working on this.

Since Sep 1, 2026.

enhancement
Dominant language
Python
Stars
3.6k
Forks
487
Avg merge
1d 1h
Merged PRs (30d)
207

Description

I'm proposing TransformerModule: the transformer parallel of the LSTMModule, built on the same four aspects. Collection runs one step at a time against a fixed-shape KV cache carried per env; training processes full windows under a block-diagonal causal mask so attention never crosses episode boundaries. The load-bearing invariant, enforced by a parity test, is that the two paths produce identical outputs. One deliberate rule throughout: the KV cache is inference state, not experience. It lives in the rollout, never enters replay, and training recomputes state under current weights.

I am self-assigning this feature to myself. I have some experiments running for this

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.