Skip to content

Conversation

@vmoens
Copy link
Collaborator

@vmoens vmoens commented Feb 3, 2023

Description

Stacking specs is a mileston in allowing for multi-task and multi-agent env wrappers.

This PR introduces two classes:

  • LazyStackedTensorSpec: represents a stack of tensorspecs. It does not require them to be homogeneous in shape but we don't support different spec type.
  • LazyStackedCompositeSpec, which inherits from the one above, but has extra key-features. Its (random) samples are of the LazyStackedTensorDict type.

The supported operations are:

  • clone
  • expand
  • to
  • rand
  • zero

Indexing can be performed along the stack dimension only. If the indexing is attempted along another dimension, an error will be thrown. This means that a stack of a stack can be indexed but only in the right order.

TODO:

  • Document the feature
  • Compatibility with nested tensors

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Feb 3, 2023
@vmoens vmoens added the enhancement New feature or request label Feb 6, 2023
@vmoens vmoens mentioned this pull request Feb 7, 2023
5 tasks
@vmoens vmoens changed the title [WIP,Feature] Stacking specs [Feature] Stacking specs Mar 15, 2023
@vmoens vmoens merged commit 3dd176d into main Mar 17, 2023
@vmoens vmoens deleted the stack_specs branch March 17, 2023 08:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants