Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions embodichain/data/enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,14 @@ class SemanticMask(IntEnum):
Attributes:
BACKGROUND (int): Represents the background region (value: 0).
FOREGROUND (int): Represents the foreground objects (value: 1).
ROBOT (int): Represents the robot region (value: 2).
ROBOT_LEFT (int): Represents the left robot region (value: 2).
ROBOT_RIGHT (int): Represents the right robot region (value: 3).
"""

BACKGROUND = 0
FOREGROUND = 1
ROBOT = 2
ROBOT_LEFT = 2
ROBOT_RIGHT = 3
Comment on lines 31 to +34
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

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

Renaming/removing SemanticMask.ROBOT is a breaking API change for any downstream code that imports that enum member or expects label value 2 to mean "robot" in general. If backward compatibility matters, consider keeping ROBOT = ROBOT_LEFT as an alias (or deprecating it) while introducing ROBOT_LEFT/ROBOT_RIGHT, and document how consumers should migrate.

Copilot uses AI. Check for mistakes.


class EndEffector(Enum):
Expand Down
37 changes: 29 additions & 8 deletions embodichain/lab/gym/envs/managers/observations.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,10 +180,11 @@ def compute_semantic_mask(
"""Compute the semantic mask for the specified scene entity.

Note:
The semantic mask is defined as (B, H, W, 3) where the three channels represents:
- robot channel: the instance id of the robot is set to 1 (0 if not robot)
The semantic mask is defined as (B, H, W, len(SemanticMask)) where these channels represents:
- background channel: the instance id of the background is set to 1 (0 if not background)
- foreground channel: the instance id of the foreground objects is set to 1 (0 if not foreground)
- robot left-side channel: the instance id of the robot left-side is set to 1
- robot right-side channel: the instance id of the robot right-side is set to 1

Args:
env: The environment instance.
Expand All @@ -209,13 +210,30 @@ def compute_semantic_mask(
else:
mask = obs["sensor"][entity_cfg.uid]["mask"]

robot_uids = env.robot.get_user_ids()
left_robot_uids = torch.cat(
[
env.robot.get_user_ids(link_name)
for link_name in env.robot.link_names
if link_name.startswith("left_")
],
-1,
)
right_robot_uids = torch.cat(
[
env.robot.get_user_ids(link_name)
for link_name in env.robot.link_names
if link_name.startswith("right_")
],
-1,
)

Comment on lines +213 to 229
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

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

torch.cat([...]) will raise at runtime if env.robot.link_names has no entries starting with left_ or right_ (empty list). To make this robust across robot models, consider getting all_robot_uids = env.robot.get_user_ids() once and slicing by the left/right link indices; if an index list is empty, use an empty tensor with the right shape/device so the mask simply becomes all-false instead of crashing.

Suggested change
left_robot_uids = torch.cat(
[
env.robot.get_user_ids(link_name)
for link_name in env.robot.link_names
if link_name.startswith("left_")
],
-1,
)
right_robot_uids = torch.cat(
[
env.robot.get_user_ids(link_name)
for link_name in env.robot.link_names
if link_name.startswith("right_")
],
-1,
)
# Get all robot user IDs once and slice by left/right link indices.
all_robot_uids = env.robot.get_user_ids()
left_link_indices = [
i
for i, link_name in enumerate(env.robot.link_names)
if link_name.startswith("left_")
]
right_link_indices = [
i
for i, link_name in enumerate(env.robot.link_names)
if link_name.startswith("right_")
]
if left_link_indices:
left_robot_uids = all_robot_uids[left_link_indices]
else:
left_robot_uids = all_robot_uids.new_empty((0,))
if right_link_indices:
right_robot_uids = all_robot_uids[right_link_indices]
else:
right_robot_uids = all_robot_uids.new_empty((0,))

Copilot uses AI. Check for mistakes.
mask_exp = mask.unsqueeze(-1)

robot_uids_exp = robot_uids.unsqueeze_(1).unsqueeze_(1)
left_robot_uids_exp = left_robot_uids.unsqueeze_(1).unsqueeze_(1)
right_robot_uids_exp = right_robot_uids.unsqueeze_(1).unsqueeze_(1)

robot_mask = (mask_exp == robot_uids_exp).any(-1).squeeze_(-1)
left_robot_mask = (mask_exp == left_robot_uids_exp).any(-1).squeeze_(-1)
right_robot_mask = (mask_exp == right_robot_uids_exp).any(-1).squeeze_(-1)

asset_uids = env.sim.asset_uids
foreground_assets = [
Expand All @@ -239,17 +257,20 @@ def compute_semantic_mask(

foreground_mask = (mask_exp == foreground_uids_exp).any(-1).squeeze_(-1)

background_mask = ~(robot_mask | foreground_mask).squeeze_(-1)
background_mask = ~(left_robot_mask | right_robot_mask | foreground_mask).squeeze_(
-1
)

masks = [None, None, None]
masks = [None, None, None, None]
masks_ids = [member.value for member in SemanticMask]
assert len(masks) == len(
masks_ids
), "Different length of mask slots and SemanticMask Enum {}.".format(masks_ids)
mask_id_to_label = {
SemanticMask.BACKGROUND.value: background_mask,
SemanticMask.FOREGROUND.value: foreground_mask,
SemanticMask.ROBOT.value: robot_mask,
SemanticMask.ROBOT_LEFT.value: left_robot_mask,
SemanticMask.ROBOT_RIGHT.value: right_robot_mask,
}
for mask_id in masks_ids:
masks[mask_id] = mask_id_to_label[mask_id]
Comment on lines 263 to 276
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

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

The semantic mask documentation and construction are now out of sync with the behavior: this code returns 4 channels (background/foreground/robot_left/robot_right), but the docstring above still describes 3 channels and the return shape as (num_envs, height, width). Also, hard-coding masks = [None, None, None, None] is brittle—prefer allocating based on len(SemanticMask) to avoid future enum changes silently breaking this function.

Copilot uses AI. Check for mistakes.
Expand Down