Skip to content

[mujoco] add native frame_stack across Gym, DMC, and robotics#380

Merged
Trinkle23897 merged 3 commits into
mainfrom
jiayi/issue136-framestack
Apr 5, 2026
Merged

[mujoco] add native frame_stack across Gym, DMC, and robotics#380
Trinkle23897 merged 3 commits into
mainfrom
jiayi/issue136-framestack

Conversation

@Trinkle23897
Copy link
Copy Markdown
Collaborator

Summary

  • Problem: MuJoCo frame stacking was missing from the native env observation path, so users needed Python-side wrappers and DMC/robotics did not preserve native multi-leaf observation structure.
  • Scope: Adds native C++ frame_stack support for Gym MuJoCo, dm_control suite envs, and Gymnasium-Robotics envs. This does not add pixel observations to the main obs pipeline.
  • Outcome: frame_stack > 1 prepends a stack dimension per observation leaf, reset fills the stack with the first frame, step rolls and appends the latest frame, and frame_stack = 1 preserves existing shapes and values.

This diff adds native frame stacking for MuJoCo state observations across Gym, DMC, and robotics env families.

Technical Details

  • Approach: Add a shared MuJoCo frame-stack helper that keeps one rolling buffer per observation key, then wire each env family to call it from C++ WriteState.
  • Code pointers:
    • envpool/mujoco/frame_stack.h: shared StackSpec and per-key FrameStackBuffer.
    • envpool/mujoco/gym/mujoco_env.h: Gym MuJoCo single-observation frame-stack path.
    • envpool/mujoco/dmc/mujoco_env.h: DMC multi-leaf observation frame-stack plumbing.
    • envpool/mujoco/robotics/mujoco_env.h: robotics multi-leaf observation frame-stack plumbing.
    • envpool/mujoco/robotics/kitchen.h: per-goal-leaf stack handling for Franka Kitchen.
  • Notes: frame_stack = 1 is intentionally a no-op in both spec shape and write-path behavior.

Test Plan

Automated

  • git diff --check: passed.
  • bazel test //envpool/mujoco:mujoco_envpool_test: passed.
  • bazel test //envpool/mujoco:mujoco_gym_deterministic_test --test_filter=frame_stack: passed.
  • bazel test //envpool/mujoco:mujoco_dmc_suite_deterministic_test --test_filter=frame_stack: passed.
  • bazel test //envpool/mujoco:robotics_test --test_filter=frame_stack: passed.

Suggested Manual

  • python - <<'PY'
    from envpool.registration import make_gymnasium
    env = make_gymnasium("HalfCheetah-v5", num_envs=1, frame_stack=4)
    obs, _ = env.reset()
    print(obs.shape)
    PY
    Verify a simple Gym MuJoCo env reports the expected stacked observation shape.

@Trinkle23897 Trinkle23897 merged commit 487ab61 into main Apr 5, 2026
13 of 14 checks passed
@Trinkle23897 Trinkle23897 deleted the jiayi/issue136-framestack branch April 5, 2026 16:54
@Trinkle23897 Trinkle23897 linked an issue Apr 5, 2026 that may be closed by this pull request
9 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature Request] Add Pixel-level Observations for mujoco

1 participant