Fix semantic mask robot user_ids indexing#222
Conversation
Build left/right link index lists first, then index robot.user_ids in a single tensor operation instead of calling get_user_ids per link and concatenating. This fixes incorrect uid gathering and is more efficient. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
This PR fixes incorrect per-environment robot user-ID aggregation in compute_semantic_mask, ensuring left/right robot link user IDs are gathered and compared against the segmentation mask with the correct tensor shape semantics.
Changes:
- Replaces per-link
get_user_ids(link_name)concatenation (which flattened across envs) with per-link index collection viaenumerate. - Indexes
env.robot.user_ids[:, link_indices]to keep user IDs in(num_envs, num_links_side)form for correct broadcasting/masking.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| left_link_indices = [ | ||
| i | ||
| for i, link_name in enumerate(env.robot.link_names) | ||
| if link_name.startswith("left_") | ||
| ] | ||
| left_robot_uids = env.robot.user_ids[:, left_link_indices] |
There was a problem hiding this comment.
The updated left/right user ID indexing fixes a per-environment shape issue, but there is no unit test covering compute_semantic_mask. Please add a test (in tests/gym/envs/managers/test_observation_functors.py, where other observation functors are tested) that uses >1 env with distinct robot.user_ids per env and verifies the resulting left/right robot mask channels are computed per-env (not flattened/concatenated across envs).
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Description
This PR fixes the
compute_semantic_maskfunction in the observation manager. The previous implementation calledget_user_ids(link_name)per link in a list comprehension and concatenated the results, which produced incorrect user ID gathering. The fix builds left/right link index lists first viaenumerate, then indexesrobot.user_ids[:, link_indices]in a single tensor operation.Type of change
Checklist
black .command to format the code base.🤖 Generated with Claude Code