Skip to content

Add planner grid cell sampler#147

Merged
yuecideng merged 3 commits intomainfrom
yueci/add-spatial-sampler
Feb 25, 2026
Merged

Add planner grid cell sampler#147
yuecideng merged 3 commits intomainfrom
yueci/add-spatial-sampler

Conversation

@yuecideng
Copy link
Contributor

@yuecideng yuecideng commented Feb 25, 2026

Description

Add a spatial functor to perform spatial planner placement sampling.

Type of change

  • New feature (non-breaking change which adds functionality)
  • Documentation update

Checklist

  • I have run the 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.

Copilot AI review requested due to automatic review settings February 25, 2026 09:33
@yuecideng yuecideng added the gym robot learning env and its related features label Feb 25, 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

This PR adds a new spatial event functor to sample planner placement positions from a 2D grid, and updates the event-functors documentation to include newly available functors (including the new sampler).

Changes:

  • Add planner_grid_cell_sampler functor for grid-based object placement sampling.
  • Update event functor docs to list additional physics/visual/spatial/asset functors, including the new sampler.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 6 comments.

File Description
embodichain/lab/gym/envs/managers/randomization/spatial.py Introduces planner_grid_cell_sampler functor for sampling grid cells and placing rigid objects.
docs/source/overview/gym/event_functors.md Documents additional event functors and adds an entry for planner_grid_cell_sampler.

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

Comment on lines +507 to +509
obj_positions[obj_id][env_id] = torch.tensor(
[x, y, z], dtype=torch.float32, device=env.device
)
Copy link

Copilot AI Feb 25, 2026

Choose a reason for hiding this comment

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

obj_positions[obj_id] is shaped (len(env_ids), 3), but it’s indexed with env_id (the global env index). If env_ids isn’t 0..len(env_ids)-1 (e.g., [2, 5]), this will be out-of-bounds or write the wrong row. Index by the loop-local row (e.g., enumerate env_ids) or build an env_id→row mapping.

Copilot uses AI. Check for mistakes.
elif not isinstance(env_ids, torch.Tensor):
env_ids = torch.tensor(env_ids, device=env.device)

self.reset(env_ids)
Copy link

Copilot AI Feb 25, 2026

Choose a reason for hiding this comment

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

The docstring says cells “will not be resampled until the grid is reset”, but __call__ unconditionally calls self.reset(env_ids) here, clearing occupancy on every invocation. If the intent is persistent no-replacement across resets, remove this call and let the EventManager/caller trigger reset() explicitly.

Suggested change
self.reset(env_ids)

Copilot uses AI. Check for mistakes.
Comment on lines +484 to +489
if len(available_cells[0]) == 0:
logger.log_warning(
f"No available cells in grid for environment {env_id_int}. All cells occupied."
)
break

Copy link

Copilot AI Feb 25, 2026

Choose a reason for hiding this comment

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

If the grid has no available cells, the code breaks out of the per-object loop, leaving later objects in that env without a sampled position. Because obj_positions is initialized to zeros, those unplaced objects will later get a zero position applied (teleporting to origin). Track placement success per object/env and skip pose updates for unplaced entries (or preserve current pose).

Copilot uses AI. Check for mistakes.
self._grid_state[env_id_int] = torch.zeros(
rows, cols, device=env.device, dtype=torch.uint8
)
self._grid_cell_sizes[env_id_int] = (cell_width, cell_height)
Copy link

Copilot AI Feb 25, 2026

Choose a reason for hiding this comment

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

self._grid_cell_sizes is written (self._grid_cell_sizes[env_id_int] = ...) but never read anywhere in this functor. If it isn’t needed, remove it; if it is intended for downstream logic, add the corresponding usage/exposure so the extra state is justified.

Suggested change
self._grid_cell_sizes[env_id_int] = (cell_width, cell_height)

Copilot uses AI. Check for mistakes.
Comment on lines +513 to +515
pose = rigid_object.get_local_pose()[env_ids]
pose[:, 0:3] = obj_positions[obj_id]

Copy link

Copilot AI Feb 25, 2026

Choose a reason for hiding this comment

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

This assignment applies obj_positions[obj_id] to all env_ids for every object, even if some environments didn’t successfully sample a cell for that object (e.g., grid full and loop broke early). That can overwrite poses with default zeros. Consider masking the pose update per-env based on which placements succeeded.

Copilot uses AI. Check for mistakes.
@yuecideng yuecideng merged commit 4bdb21e into main Feb 25, 2026
9 checks passed
@yuecideng yuecideng deleted the yueci/add-spatial-sampler branch February 25, 2026 09:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

gym robot learning env and its related features

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants