[Refactor] TypedTensorDict redesign (#1662)

Author: vmoens

docs/source/compatibility.rst

@@ -15,7 +15,7 @@ Architecture overview
    TensorCollection
    ├── TensorDictBase
    │   ├── TensorDict                 (in-memory)
-   │   │   └── TypedTensorDict        (typed fields, IS-A TensorDict)
+   │   ├── TypedTensorDict            (typed fields, wraps any TensorDictBase)
    │   ├── PersistentTensorDict       (HDF5-backed)
    │   ├── TensorDictStore            (Redis / Dragonfly / KeyDB)
    │   └── LazyStackedTensorDict      (lazy stack of heterogeneous TDs)
@@ -26,8 +26,11 @@ Two patterns exist for adding typed field declarations:
 
 - **TensorClass** wraps any ``TensorDictBase`` via ``from_tensordict(td)``.
   It delegates all storage to the wrapped object.
-- **TypedTensorDict** *is* a ``TensorDict``.  It stores data in-memory and
-  interoperates with other backends through conversion or stacking.
+- **TypedTensorDict** wraps any ``TensorDictBase`` via ``from_tensordict(td)``,
+  similar to ``TensorClass``.  Direct construction creates a ``TensorDict``
+  internally.  Unlike ``TensorClass``, it inherits from ``TensorDictBase``
+  directly, supports ``**state`` spreading natively, and uses standard
+  Python inheritance for schema composition.
 
 TensorClass + backends
 ----------------------
@@ -205,109 +208,162 @@ enforce schemas, but they compose without conflict:
 TypedTensorDict + backends
 --------------------------
 
-``TypedTensorDict`` is a ``TensorDict`` subclass.  It stores data in-memory
-but interoperates with other backends through conversion or stacking.
+``TypedTensorDict.from_tensordict(td)`` accepts any ``TensorDictBase`` subclass,
+just like ``TensorClass``.  The backend is stored live (no copy) -- mutations
+through the ``TypedTensorDict`` go directly to the underlying backend.
+
+.. code-block:: python
+
+   from tensordict import TypedTensorDict
+   from torch import Tensor
+
+   class State(TypedTensorDict):
+       x: Tensor
+       y: Tensor
+
+   state = State.from_tensordict(some_backend)
 
 .. list-table::
    :header-rows: 1
-   :widths: 30 12 12 12 12 12
+   :widths: 22 10 10 10 10 10 10 10 10
 
-   * - Pattern
+   * - Backend
      - Build
      - Read
      - Write
      - Index
+     - Clone
      - Stack
-   * - Direct construction
-     - yes
+     - Iter
+     - Update
+   * - ``TensorDict``
      - yes
      - yes
      - yes
      - yes
-   * - From H5 (materialise then construct)
      - yes
      - yes
      - yes
      - yes
+   * - ``PersistentTensorDict`` (H5)
      - yes
-   * - From Redis (materialise then construct)
      - yes
      - yes
      - yes
      - yes
      - yes
-   * - From lazy stack (materialise then construct)
      - yes
      - yes
+   * - ``TensorDictStore`` (Redis)
      - yes
      - yes
      - yes
-   * - ``torch.stack`` (dense)
      - yes
      - yes
      - yes
      - yes
-     - --
-   * - ``LazyStackedTensorDict`` of TTDs
      - yes
+   * - ``LazyStackedTensorDict``
      - yes
      - yes
      - yes
-     - --
-   * - ``memmap_()``
      - yes
      - yes
-     - set\_()
      - yes
      - yes
-   * - To H5 (``PersistentTensorDict.from_dict``)
      - yes
+   * - ``TensorDict`` (memmap)
      - yes
-     - H5 rules
      - yes
-     - --
-   * - To Redis (``TensorDictStore.from_tensordict``)
+     - set\_()
      - yes
      - yes
      - yes
      - yes
-     - --
+     - update\_()
 
-Constructing TypedTensorDict from other backends
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. note::
 
-Since ``TypedTensorDict`` is an in-memory ``TensorDict``, loading data from a
-remote or persistent backend requires materialising the data first:
+   Memory-mapped TensorDicts are locked after ``memmap_()``.  Use
+   ``set_()`` and ``update_()`` for in-place writes instead of attribute
+   assignment or ``update()``.
+
+Building a TypedTensorDict on each backend
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+**In-memory TensorDict** -- the default (direct construction creates one
+internally):
 
 .. code-block:: python
 
    >>> import torch
-   >>> from tensordict import TypedTensorDict
+   >>> from tensordict import TensorDict, TypedTensorDict
    >>> from torch import Tensor
    >>>
    >>> class State(TypedTensorDict):
    ...     x: Tensor
    ...     y: Tensor
+   >>>
+   >>> state = State(x=torch.randn(4, 3), y=torch.randn(4, 5), batch_size=[4])
+   >>> state.x.shape
+   torch.Size([4, 3])
 
-**From HDF5**:
+**Wrapping an existing TensorDict** via ``from_tensordict`` (zero-copy):
+
+.. code-block:: python
+
+   >>> td = TensorDict(x=torch.randn(4, 3), y=torch.randn(4, 5), batch_size=[4])
+   >>> state = State.from_tensordict(td)
+   >>> state.x.shape  # reads from td
+   torch.Size([4, 3])
+   >>> state.x = torch.ones(4, 3)  # writes to td
+   >>> (td["x"] == 1).all()
+   True
+
+**HDF5 (PersistentTensorDict)**:
 
 .. code-block:: python
 
    >>> from tensordict import PersistentTensorDict
    >>>
    >>> h5 = PersistentTensorDict.from_h5("data.h5")
-   >>> local = h5.to_tensordict()
-   >>> state = State(x=local["x"], y=local["y"], batch_size=local.batch_size)
+   >>> state = State.from_tensordict(h5)
+   >>> state.x.shape  # reads from HDF5
+   torch.Size([4, 3])
 
-**From a lazy stack**:
+**Redis (TensorDictStore)**:
+
+.. code-block:: python
+
+   >>> from tensordict.store import TensorDictStore
+   >>>
+   >>> store = TensorDictStore.from_tensordict(td, host="localhost")
+   >>> state = State.from_tensordict(store)
+   >>> state.x.shape  # fetched from Redis
+   torch.Size([4, 3])
+
+**Lazy stack**:
 
 .. code-block:: python
 
    >>> from tensordict import lazy_stack
    >>>
-   >>> ls = lazy_stack([td1, td2], dim=0)
-   >>> local = ls.to_tensordict()
-   >>> state = State(x=local["x"], y=local["y"], batch_size=local.batch_size)
+   >>> tds = [TensorDict(x=torch.randn(3), y=torch.randn(5)) for _ in range(4)]
+   >>> ls = lazy_stack(tds, dim=0)
+   >>> state = State.from_tensordict(ls)
+   >>> state[0].x.shape
+   torch.Size([3])
+
+**Memory-mapped TensorDict**:
+
+.. code-block:: python
+
+   >>> td_mmap = td.memmap_("/tmp/my_memmap")
+   >>> state = State.from_tensordict(td_mmap)
+   >>> state.x.shape
+   torch.Size([4, 3])
+   >>> # memmap TDs are locked -- use in-place operations:
+   >>> state.set_("x", torch.ones(4, 3))
 
 Stacking TypedTensorDicts
 ^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -334,30 +390,69 @@ Lazy stacking also works.  Indexing a ``LazyStackedTensorDict`` of
    >>> isinstance(ls[0], State)
    True
 
-Saving TypedTensorDict to persistent backends
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. _compat-redis-prealloc:
+
+Pre-allocating on Redis and filling iteratively
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Since ``TypedTensorDict`` is a ``TensorDict``, it can be saved to HDF5, Redis,
-or memory-mapped storage directly:
+A common pattern for shared replay buffers or distributed data stores is to
+pre-allocate storage on a remote server (Redis / Dragonfly / KeyDB) and fill
+it one sample at a time, without ever loading the full dataset into RAM.
+
+``TensorDictStore.from_schema`` creates keys with known shapes and dtypes
+directly on the server using ``SETRANGE`` (zero-filled by the server; no
+tensor data passes through Python):
 
 .. code-block:: python
 
-   >>> # To HDF5
-   >>> from tensordict import PersistentTensorDict
-   >>> h5 = PersistentTensorDict.from_dict(state, filename="state.h5")
+   >>> import torch
+   >>> from tensordict import TensorDict, TypedTensorDict
+   >>> from tensordict.store import TensorDictStore
+   >>> from torch import Tensor
    >>>
-   >>> # To memmap
-   >>> state.memmap_("/tmp/state_mmap")
+   >>> class Replay(TypedTensorDict):
+   ...     obs: Tensor
+   ...     action: Tensor
+   ...     reward: Tensor
    >>>
-   >>> # To Redis
-   >>> from tensordict.store import TensorDictStore
-   >>> store = TensorDictStore.from_tensordict(state, host="localhost")
+   >>> # Pre-allocate 100k entries directly on Redis -- no RAM used
+   >>> store = TensorDictStore.from_schema(
+   ...     {"obs": ([84, 84, 3], torch.uint8),
+   ...      "action": ([4], torch.float32),
+   ...      "reward": ([], torch.float32)},
+   ...     batch_size=[100_000],
+   ...     host="redis-node",
+   ... )
+   >>>
+   >>> # Wrap with typed access
+   >>> replay = Replay.from_tensordict(store)
+   >>>
+   >>> # Fill iteratively -- each write goes directly to Redis
+   >>> for i, sample in enumerate(data_stream):
+   ...     replay[i] = Replay(
+   ...         obs=sample.obs, action=sample.action, reward=sample.reward,
+   ...         batch_size=[],
+   ...     )
+
+If the store is initially empty (no keys registered yet), use ``check=False``
+to skip the key-presence validation and fill keys on the fly:
+
+.. code-block:: python
+
+   >>> store = TensorDictStore(batch_size=[100_000], host="redis-node")
+   >>> replay = Replay.from_tensordict(store, check=False)
+   >>>
+   >>> # First indexed write auto-creates each key via SETRANGE
+   >>> replay[0] = Replay(obs=obs_0, action=act_0, reward=r_0, batch_size=[])
+   >>> # Subsequent writes fill in the pre-allocated storage
+   >>> replay[1] = Replay(obs=obs_1, action=act_1, reward=r_1, batch_size=[])
 
 
 TensorClass vs TypedTensorDict
 ------------------------------
 
-Both enforce typed schemas but differ architecturally:
+Both enforce typed schemas and can wrap any ``TensorDictBase`` backend, but
+they differ architecturally:
 
 .. list-table::
    :header-rows: 1
@@ -366,12 +461,12 @@ Both enforce typed schemas but differ architecturally:
    * - Aspect
      - ``TensorClass``
      - ``TypedTensorDict``
-   * - Relationship to ``TensorDict``
-     - Wraps a ``TensorDictBase`` (HAS-A)
-     - Is a ``TensorDict`` (IS-A)
+   * - Relationship to ``TensorDictBase``
+     - Wraps a ``TensorDictBase`` (HAS-A via ``TensorCollection``)
+     - Is a ``TensorDictBase`` (IS-A, delegates to ``_source``)
    * - Can wrap non-TensorDict backends
      - Yes (H5, Redis, lazy stack, etc.)
-     - No (in-memory only; convert first)
+     - Yes (H5, Redis, lazy stack, etc.)
    * - ``**state`` spreading
      - Field-by-field repacking
      - Natively (``MutableMapping``)
@@ -380,15 +475,19 @@ Both enforce typed schemas but differ architecturally:
      - Not supported (tensor-only)
    * - Backend stays live
      - Yes (writes go to original backend)
-     - No (data is in-memory after construction)
+     - Yes (writes go to original backend)
+   * - Python inheritance
+     - Not supported
+     - Supported (standard class hierarchy)
    * - Composable with each other
      - Yes (``TC.from_tensordict(ttd)`` works)
-     - N/A
+     - Yes (``TTD.from_tensordict(tc._tensordict)`` works)
 
-When a ``TensorClass`` wraps a persistent backend (H5, Redis), writes through
-the ``TensorClass`` go directly to that backend.  When a ``TypedTensorDict`` is
-constructed from persistent data, the data is copied into memory.
+Both wrappers keep the backend alive -- mutations through the typed wrapper go
+directly to the underlying storage.  Direct construction (without
+``from_tensordict``) creates an in-memory ``TensorDict`` as the backend.
 
-Choose ``TensorClass`` when you need live access to a remote or on-disk backend
-with typed field access.  Choose ``TypedTensorDict`` when you want typed,
-in-memory state with ``**state`` spreading and standard Python inheritance.
+Choose ``TensorClass`` when you need non-tensor fields or want to integrate
+with existing tensorclass-based APIs.  Choose ``TypedTensorDict`` when you
+want native ``**state`` spreading, standard Python inheritance for schema
+composition, and full ``TensorDictBase`` API compatibility.