Fix issue in obs manager for adding new key#169
Conversation
There was a problem hiding this comment.
Pull request overview
This PR refactors assign_data_to_dict in embodichain/lab/gym/utils/gym_utils.py to operate exclusively on TensorDict, aligning observation assignment with the codebase’s TensorDict-based observation structures.
Changes:
- Updates
assign_data_to_dictsignature to acceptTensorDictinstead of supporting bothdictandTensorDict. - Creates missing intermediate nested containers as
TensorDictinstances inheritingbatch_sizeanddevicefrom the parent. - Uses
TensorDict.get()for type checking to avoid side effects during existence checks.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 1 out of 1 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| last_key = keys[-1] | ||
| current_data.batch_size = batch_size # Ensure the batch size is consistent | ||
| current_data[last_key] = value |
There was a problem hiding this comment.
current_data.batch_size = batch_size is likely invalid for tensordict.TensorDict (batch_size is typically derived/validated and may be read-only), and even if it works it can silently desynchronize the nested TensorDict from the shapes of its contained tensors. Prefer enforcing batch size by creating intermediate TensorDicts with the intended batch_size (e.g., from the root) and, when traversing existing nested entries, validate their batch_size matches instead of mutating it.
| def assign_data_to_dict(data_dict: TensorDict, name: str, value: Any) -> None: | ||
| """Assign data to a TensorDict using a '/' separated key. | ||
| Missing intermediate TensorDicts will be created automatically. | ||
|
|
||
| Args: | ||
| data_dict (Dict[str, Union[Any, Dict[str, Any]]]): The nested dictionary to assign data to. | ||
| data_dict (TensorDict): The TensorDict to assign data to. | ||
| name (str): The '/' separated key string. | ||
| value (Any): The value to assign. | ||
| """ |
There was a problem hiding this comment.
This helper now enforces TensorDict for assignment, but the paired fetch_data_from_dict helper (used in ObservationManager) is still typed/documented as operating on Dict[...]. Consider updating the fetch helper’s type hints/docstring (and any referenced docs) to TensorDict/EnvObs too, so callers don’t get misleading typing and the TensorDict-only contract is consistent.
Description
Refactor
assign_data_to_dictto only support TensorDictThis PR modifies the
assign_data_to_dictfunction inembodichain/lab/gym/utils/gym_utils.pyto exclusively support TensorDict instead of supporting both regular dict and TensorDict.Changes made:
TensorDictinstead ofDict[str, Union[Any, Dict[str, Any]]]current_data.get(key)instead ofcurrent_data[key]to avoid key creation during existence checkTensorDictwith the samebatch_sizeanddeviceas the parent when missing keys are encounteredThis change ensures consistency with other functions in the codebase that work with Tensordict and simplifies the implementation by removing dual-type support.
Type of change
Checklist
black .command to format the code base.