feat(processor): Introduce ToBatchProcessor for handling observation batching

- Added ToBatchProcessor to ensure observations have proper batch dimensions for model processing.
- Implemented functionality to add batch dimensions to state and image observations as needed.
- Created comprehensive unit tests to validate the processor's behavior with various tensor dimensions and types.
- Ensured compatibility with existing transition keys and maintained the integrity of non-observation data.
This commit is contained in:
Adil Zouitine
2025-07-23 18:41:53 +02:00
committed by Steven Palma
parent e7be2fd113
commit f5c6b03b61
3 changed files with 517 additions and 0 deletions

View File

@@ -14,6 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from .batch_processor import ToBatchProcessor
from .device_processor import DeviceProcessor
from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor
from .observation_processor import VanillaObservationProcessor
@@ -48,6 +49,7 @@ __all__ = [
"RenameProcessor",
"RewardProcessor",
"RobotProcessor",
"ToBatchProcessor",
"TransitionKey",
"TruncatedProcessor",
"VanillaObservationProcessor",