-
Notifications
You must be signed in to change notification settings - Fork 2.7k
Add APIs to Isaac Lab Mimic for supporting loco-manipulation datagen #3992
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
faaed31
a306cbb
64feab8
3ae7ed9
02a2670
b03090a
d057460
24f9b80
3968c5f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -11,6 +11,12 @@ | |||||
| from isaaclab.envs import ManagerBasedRLEnv | ||||||
|
|
||||||
|
|
||||||
| def optional_method(func): | ||||||
| """Decorator to mark a method as optional.""" | ||||||
| func.__is_optional__ = True | ||||||
| return func | ||||||
|
|
||||||
|
|
||||||
| class ManagerBasedRLMimicEnv(ManagerBasedRLEnv): | ||||||
| """The superclass for the Isaac Lab Mimic environments. | ||||||
|
|
||||||
|
|
@@ -156,3 +162,21 @@ def serialize(self): | |||||
| and used in utils/env_utils.py. | ||||||
| """ | ||||||
| return dict(env_name=self.spec.id, type=2, env_kwargs=dict()) | ||||||
|
|
||||||
| @optional_method | ||||||
| def get_navigation_state(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: | ||||||
| """ | ||||||
| Optional method. Only required when using navigation controller locomanipulation data generation. | ||||||
|
|
||||||
| Gets the navigation state of the robot. Required when use of the navigation controller is | ||||||
| enabled. The navigation state includes a boolean flag "is_navigating" to indicate when the | ||||||
| robot is under control by the navigation controller, and a boolean flag "navigation_goal_reached" | ||||||
| to indicate when the navigation goal has been reached. | ||||||
|
|
||||||
| Args: | ||||||
| env_id: The environment index to get the navigation state for. If None, all envs are considered. | ||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. syntax: docstring parameter name mismatch: documented as
Suggested change
|
||||||
|
|
||||||
| Returns: | ||||||
| A dictionary that of navigation state flags (False or True). | ||||||
| """ | ||||||
| raise NotImplementedError | ||||||
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -7,10 +7,13 @@ | |||||
| Base class for data generator. | ||||||
| """ | ||||||
| import asyncio | ||||||
| import copy | ||||||
| import numpy as np | ||||||
| import torch | ||||||
| from typing import Any | ||||||
|
|
||||||
| import omni.log | ||||||
|
|
||||||
| import isaaclab.utils.math as PoseUtils | ||||||
| from isaaclab.envs import ( | ||||||
| ManagerBasedRLMimicEnv, | ||||||
|
|
@@ -688,6 +691,10 @@ async def generate( # noqa: C901 | |||||
| eef_subtasks_done[eef_name] = False | ||||||
|
|
||||||
| prev_src_demo_datagen_info_pool_size = 0 | ||||||
|
|
||||||
| if self.env_cfg.datagen_config.use_navigation_controller: | ||||||
| was_navigating = False | ||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. logic: |
||||||
|
|
||||||
| # While loop that runs per time step | ||||||
| while True: | ||||||
| async with self.src_demo_datagen_info_pool.asyncio_lock: | ||||||
|
|
@@ -880,8 +887,54 @@ async def generate( # noqa: C901 | |||||
| generated_actions.extend(exec_results["actions"]) | ||||||
| generated_success = generated_success or exec_results["success"] | ||||||
|
|
||||||
| # Get the navigation state | ||||||
| if self.env_cfg.datagen_config.use_navigation_controller: | ||||||
| processed_nav_subtask = False | ||||||
| navigation_state = self.env.get_navigation_state(env_id) | ||||||
| assert navigation_state is not None, "Navigation state cannot be None when using navigation controller" | ||||||
| is_navigating = navigation_state["is_navigating"] | ||||||
| navigation_goal_reached = navigation_state["navigation_goal_reached"] | ||||||
peterd-NV marked this conversation as resolved.
Show resolved
Hide resolved
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. logic: potential KeyError if
Suggested change
|
||||||
|
|
||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. logic: potential KeyError if
Suggested change
|
||||||
| for eef_name in self.env_cfg.subtask_configs.keys(): | ||||||
| current_eef_subtask_step_indices[eef_name] += 1 | ||||||
|
|
||||||
| # Execute locomanip navigation controller if it is enabled via the use_navigation_controller flag | ||||||
| if self.env_cfg.datagen_config.use_navigation_controller: | ||||||
| if "body" not in self.env_cfg.subtask_configs.keys(): | ||||||
| error_msg = ( | ||||||
| 'End effector with name "body" not found in subtask configs. "body" must be a valid end' | ||||||
| " effector to use the navigation controller.\n" | ||||||
| ) | ||||||
| omni.log.error(error_msg) | ||||||
| raise RuntimeError(error_msg) | ||||||
|
Comment on lines
+903
to
+909
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. style: inefficient to validate "body" existence on every iteration of the |
||||||
|
|
||||||
| # Repeat the last nav subtask action if the robot is navigating and hasn't reached the waypoint goal | ||||||
| if ( | ||||||
| current_eef_subtask_step_indices["body"] == len(current_eef_subtask_trajectories["body"]) - 1 | ||||||
| and not processed_nav_subtask | ||||||
| ): | ||||||
| if is_navigating and not navigation_goal_reached: | ||||||
| for name in self.env_cfg.subtask_configs.keys(): | ||||||
| current_eef_subtask_step_indices[name] -= 1 | ||||||
| processed_nav_subtask = True | ||||||
|
|
||||||
| # Else skip to the end of the nav subtask if the robot has reached the waypoint goal before the end | ||||||
| # of the human recorded trajectory | ||||||
| elif was_navigating and not is_navigating and not processed_nav_subtask: | ||||||
| number_of_steps_to_skip = len(current_eef_subtask_trajectories["body"]) - ( | ||||||
| current_eef_subtask_step_indices["body"] + 1 | ||||||
| ) | ||||||
| for name in self.env_cfg.subtask_configs.keys(): | ||||||
| if current_eef_subtask_step_indices[name] + number_of_steps_to_skip < len( | ||||||
| current_eef_subtask_trajectories[name] | ||||||
| ): | ||||||
| current_eef_subtask_step_indices[name] = ( | ||||||
| current_eef_subtask_step_indices[name] + number_of_steps_to_skip | ||||||
| ) | ||||||
| else: | ||||||
| current_eef_subtask_step_indices[name] = len(current_eef_subtask_trajectories[name]) - 1 | ||||||
| processed_nav_subtask = True | ||||||
|
|
||||||
| subtask_ind = current_eef_subtask_indices[eef_name] | ||||||
| if current_eef_subtask_step_indices[eef_name] == len( | ||||||
| current_eef_subtask_trajectories[eef_name] | ||||||
|
|
@@ -923,6 +976,10 @@ async def generate( # noqa: C901 | |||||
| else: | ||||||
| current_eef_subtask_step_indices[eef_name] = None | ||||||
| current_eef_subtask_indices[eef_name] += 1 | ||||||
|
|
||||||
| if self.env_cfg.datagen_config.use_navigation_controller: | ||||||
| was_navigating = copy.deepcopy(is_navigating) | ||||||
|
|
||||||
| # Check if all eef_subtasks_done values are True | ||||||
| if all(eef_subtasks_done.values()): | ||||||
| break | ||||||
|
|
||||||
Uh oh!
There was an error while loading. Please reload this page.