[Feature Request] Make transforms stateless
@vmoens is already working on this.
Since Jun 2, 2023.
- Dominant language
- Python
- Stars
- 3.6k
- Forks
- 487
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 207
Description
Motivation
TensorDict is a great tool to carry information from one step to another.
Yet, we still use buffers in transforms like CatFrames or ObservationNorm.
We could spare these by using tensors stored in a tensordict.
The only drawback is memory consumption: we may not want to carry over these buffer in the replay buffer etc.
One way to go about that would be to have a prefix in these keys (or even better: store them in a sub-tensordict) that we could discard in env.rollout or collector.rollout.
Something like
>>> env.step()
TensorDict({
"_transform_buffers": TensorDict({
"0": TensorDict({
"loc": torch.Tensor(...)
}, batch_size=[])
}, batch_size=[])
}, batch_size=[])
Open questions:
We need to figure out how to transfer some normalization constants from one env to another.
If we do
train_env = make_env()
t = train_env.transform[0] # an ObservationNorm
t.init_stats() # Presumably: will create a TensorDictPrimer in the env with the normalization constants.
test_env()
? # here we need to figure out how to gather the stats from the other env
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.