Implementation:Haosulab ManiSkill FrameStackWrapper

Knowledge Sources	Haosulab_ManiSkill
Domains	Robotics, Simulation, Environment Wrapper
Last Updated	2026-02-15 08:00 GMT

Overview

Concrete tool for stacking consecutive observations in a rolling buffer to provide temporal context to agents in GPU-vectorized environments.

Description

The FrameStack wrapper is a Gymnasium ObservationWrapper adapted from the standard gymnasium frame stack to support GPU-vectorized ManiSkill environments with both tensor and dict observations.

Key behavior:

Maintains a deque of the most recent num_stack observations.
On reset(), fills the entire buffer with the initial observation (so the agent sees num_stack identical frames initially).
On step(), appends the new observation and returns the stacked version.
For tensor observations, stacking produces shape (B, num_stack, ...) via torch.stack().transpose(0, 1).
For dict observations, each leaf tensor is stacked independently.
Automatically updates the environment's observation space to reflect the stacked shape.

Limitations:

Partial environment reset is not supported for GPU-parallelized environments (raises RuntimeError).
The lz4 compression option is accepted but not currently used.

Usage

Wrap a ManiSkill environment when the RL agent needs temporal context (e.g., for velocity estimation from position-only observations, or for recurrent/transformer policies).

Code Reference

Source Location

Repository: Haosulab_ManiSkill
File: mani_skill/utils/wrappers/frame_stack.py

Signature

class FrameStack(gym.ObservationWrapper):
    def __init__(self, env: gym.Env, num_stack: int, lz4_compress: bool = False): ...
    def observation(self, observation) -> Union[torch.Tensor, dict]: ...
    def step(self, action) -> tuple: ...
    def reset(self, seed=None, options=None) -> tuple: ...

Import

from mani_skill.utils.wrappers.frame_stack import FrameStack

I/O Contract

Inputs

Name	Type	Required	Description
env	gym.Env	Yes	The environment to wrap
num_stack	int	Yes	Number of frames to stack
lz4_compress	bool	No	LZ4 compression flag (not actively used)

Outputs

Name	Type	Description
obs	torch.Tensor or dict	Stacked observations with shape (B, num_stack, ...)
reward	torch.Tensor	Reward from the latest step
terminated/truncated	torch.Tensor	Termination signals
info	dict	Info from latest step

Usage Examples

Basic Usage

import gymnasium as gym
from mani_skill.utils.wrappers.frame_stack import FrameStack

env = gym.make("PickCube-v1", num_envs=4, obs_mode="state")
env = FrameStack(env, num_stack=4)

obs, info = env.reset()
# obs shape: (4, 4, 42) -- 4 envs, 4 stacked frames, 42-dim state

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment