Skip to content

feat(obs): add support for extra observations from gym_config#205

Merged
yuecideng merged 1 commit intomainfrom
feat/extra-obs-support
Mar 26, 2026
Merged

feat(obs): add support for extra observations from gym_config#205
yuecideng merged 1 commit intomainfrom
feat/extra-obs-support

Conversation

@yuecideng
Copy link
Contributor

Description

This PR adds support for pre-allocating extra observations from gym_config's observation functors.

Previously, init_rollout_buffer_from_config only pre-allocated buffers for basic robot observations (qpos, qvel, qf) and sensor data. This made it difficult to add custom observations without running the environment first to discover the observation space.

Now, observation functors in add mode can specify a shape in their extra parameter, which will be used to pre-allocate buffers during rollout buffer initialization. This enables efficient buffer allocation for custom observations without requiring runtime introspection.

Example usage:

env:
  observations:
    custom_position:
      mode: "add"
      func: "get_object_position"
      params:
        entity_cfg: {uid: "target"}
      extra:
        shape: [3]  # Will create buffer with shape (num_envs, max_episode_steps, 3)

Dependencies: None

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • Enhancement (non-breaking change which improves an existing functionality)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (existing functionality will not work without user modification)
  • Documentation update

Screenshots

N/A

Checklist

  • I have run `black .` command to format the code base.
  • I have made corresponding changes to the documentation
  • I have added tests that prove my fix is effective or that my feature works
  • Dependencies have been updated, if applicable.

🤖 Generated with Claude Code

Add extra parameter to FunctorCfg to allow setting metadata like output shape
for observation functors. Extract extra observations in 'add' mode in
init_rollout_buffer_from_config to pre-allocate buffers based on shapes.

Changes:
- Add extra: dict[str, Any] to FunctorCfg
- Add extra attribute documentation to ObservationCfg
- Modify init_rollout_buffer_from_config to parse and add extra observations
- Add comprehensive unit tests in tests/gym/utils/test_gym_utils.py

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Copilot AI review requested due to automatic review settings March 26, 2026 08:18
@yuecideng yuecideng added gym robot learning env and its related features dataset labels Mar 26, 2026
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds support for pre-allocating rollout-buffer slots for custom (“extra”) observations defined in gym_config.env.observations functors, enabling buffer initialization without needing to run the environment to infer observation shapes.

Changes:

  • Extend init_rollout_buffer_from_config to parse env.observations in add mode and pre-allocate tensors when extra.shape is provided.
  • Add extra metadata field to FunctorCfg to carry auxiliary configuration (e.g., output shape).
  • Add a comprehensive test suite covering extra observation pre-allocation behavior and nested observation names.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File Description
embodichain/lab/gym/utils/gym_utils.py Parse env.observations (add mode) and pre-allocate buffers for entries that declare extra.shape, then insert them into rollout_buffer["obs"].
embodichain/lab/gym/envs/managers/cfg.py Add FunctorCfg.extra to support storing metadata like observation output shape.
tests/gym/utils/test_gym_utils.py New tests verifying buffer initialization for extra observations (various shapes, modes, nested keys) and interaction with sensor allocations.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +253 to +276
"sensor": [
{
"uid": "camera",
"width": 320,
"height": 240,
"enable_mask": True,
}
],
"env": {
"observations": {
"extra_vec": {
"mode": "add",
"extra": {"shape": [10]},
}
}
},
}

buffer = init_rollout_buffer_from_config(
config=config,
max_episode_steps=100,
batch_size=4,
state_dim=7,
device="cpu",
Copy link

Copilot AI Mar 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

test_sensor_and_extra_obs_together allocates a very large buffer (batch_size=4, max_episode_steps=100, 240x320 color + mask), which is ~245MB just for the sensor tensors. This can make the unit test suite slow and memory-hungry in CI. Consider reducing width/height and/or max_episode_steps/batch_size in this test while still asserting the same shape logic (e.g., use smaller dimensions like 32x24 and fewer steps).

Copilot uses AI. Check for mistakes.
@yuecideng yuecideng merged commit 96caffc into main Mar 26, 2026
21 of 24 checks passed
@yuecideng yuecideng deleted the feat/extra-obs-support branch March 26, 2026 09:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

dataset gym robot learning env and its related features

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants