Fix demo action shape normalization and control-part mapping#237
Merged
Conversation
Add automatic post-processing for create_demo_action_list outputs across inherited env classes to enforce action dimension consistency.\n\nWhen demo action dim is larger than action space, actions are sliced by active_joint_ids; when smaller, a clear error is raised.\n\nAlso align control-part sourcing with env config and update pour-water config control settings.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR aims to prevent demo-action shape mismatches by normalizing subclass create_demo_action_list outputs to the environment action dimension, and aligns configurable-action “control part” resolution with environment configuration (plus updates the pour-water config accordingly).
Changes:
- Wraps subclass
create_demo_action_listviaEmbodiedEnv.__init_subclass__and adds torch-based demo-action shape validation/conversion (including slicing viaactive_joint_ids). - Switches configurable-action control-part lookup to read from
env.cfg.control_parts. - Updates pour-water simple gym config control joints/control parts, and applies minor formatting tweaks.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
embodichain/lab/sim/robots/cobotmagic.py |
Formatting-only change to solver TCP array formatting. |
embodichain/lab/gym/envs/embodied_env.py |
Adds demo-action normalization utilities and auto-wrapping of subclass demo builders. |
embodichain/lab/gym/envs/action_bank/configurable_action.py |
Updates control-part source to use environment config. |
configs/gym/pour_water/gym_config_simple.json |
Updates eef joint IDs and adds env.control_parts list. |
Comments suppressed due to low confidence (1)
embodichain/lab/gym/envs/action_bank/configurable_action.py:1005
get_control_partassumesenv.cfg.control_partsis an iterable, butcontrol_partsis optional inEmbodiedEnvCfgand can beNone. As written,agent_uid in control_partswill raiseTypeErrorwhencontrol_partsisNone. Default to an empty list (or fall back to the prior metadata source) before doing membership tests and passing it todata_key_to_control_part.
def get_control_part(env, agent_uid):
control_parts = env.cfg.control_parts
if agent_uid in control_parts:
return agent_uid
else:
return data_key_to_control_part(
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if action_list is None: | ||
| return None | ||
|
|
||
| expected_dim = int(np.prod(self.action_space.shape)) |
Comment on lines
+649
to
+707
| def _normalize_demo_action_list( | ||
| self, action_list: Sequence[EnvAction] | torch.Tensor | None | ||
| ) -> Sequence[EnvAction] | torch.Tensor | None: | ||
| """Validate/convert demo action outputs to match single action-space dim.""" | ||
| if action_list is None: | ||
| return None | ||
|
|
||
| expected_dim = int(np.prod(self.action_space.shape)) | ||
|
|
||
| if isinstance(action_list, torch.Tensor): | ||
| return self._normalize_demo_action_tensor(action_list, expected_dim) | ||
|
|
||
| if not isinstance(action_list, Sequence): | ||
| raise TypeError( | ||
| "create_demo_action_list must return None, a torch.Tensor, or a sequence of actions. " | ||
| f"Got {type(action_list)}." | ||
| ) | ||
|
|
||
| normalized_action_list = [ | ||
| self._normalize_demo_action_tensor(action, expected_dim) | ||
| for action in action_list | ||
| ] | ||
| return type(action_list)(normalized_action_list) | ||
|
|
||
| def _normalize_demo_action_tensor( | ||
| self, action: EnvAction | torch.Tensor, expected_dim: int | ||
| ) -> EnvAction | torch.Tensor: | ||
| """Normalize one action tensor to the expected action dimension. | ||
|
|
||
| Conversion rule: | ||
| - If last-dim equals action-space dim, keep as-is. | ||
| - If last-dim is larger, slice with ``active_joint_ids``. | ||
| - If last-dim is smaller, raise ``ValueError``. | ||
| """ | ||
| if isinstance(action, TensorDict): | ||
| return self._normalize_demo_action_tensordict(action, expected_dim) | ||
|
|
||
| if not isinstance(action, torch.Tensor): | ||
| raise TypeError( | ||
| "Each demo action must be a torch.Tensor or TensorDict. " | ||
| f"Got {type(action)}." | ||
| ) | ||
|
|
||
| if action.ndim == 0: | ||
| raise ValueError( | ||
| "Demo action tensor must have at least one dimension with action features on the last axis." | ||
| ) | ||
|
|
||
| action_dim = int(action.shape[-1]) | ||
| if action_dim == expected_dim: | ||
| return action | ||
| if action_dim < expected_dim: | ||
| raise ValueError( | ||
| "Demo action dim is smaller than action space dim and cannot be auto-converted. " | ||
| f"Got action dim={action_dim}, expected={expected_dim}." | ||
| ) | ||
| return self._slice_action_with_active_joint_ids( | ||
| action, action_dim, expected_dim | ||
| ) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR normalizes demo action outputs to the environment action dimension and updates control-part wiring used by configurable actions and pour-water config.
Motivation and context:
Dependencies: None
Issue reference: N/A (no linked issue)
Type of change
Screenshots
N/A
Checklist
black .command to format the code base.Changes included
EmbodiedEnvso overriddencreate_demo_action_listoutputs are normalized automatically.Generated with GitHub Copilot