diff --git a/exts/StrideSim/StrideSim/__init__.py b/exts/StrideSim/StrideSim/__init__.py index 00c80da..330ff89 100644 --- a/exts/StrideSim/StrideSim/__init__.py +++ b/exts/StrideSim/StrideSim/__init__.py @@ -2,8 +2,11 @@ Python module serving as a project/extension template. """ -# Register Gym environments. -from .tasks import * +from .base_sample import * # Register UI extensions. -from .ui import * +from .quadruped import * +from .quadruped_extension import * + +# Register Gym environments. +from .tasks import * diff --git a/exts/StrideSim/StrideSim/base_sample/__init__.py b/exts/StrideSim/StrideSim/base_sample/__init__.py new file mode 100644 index 0000000..50af2f7 --- /dev/null +++ b/exts/StrideSim/StrideSim/base_sample/__init__.py @@ -0,0 +1,2 @@ +from .base_sample import BaseSample +from .base_sample_extension import BaseSampleExtension diff --git a/exts/StrideSim/StrideSim/base_sample/base_sample.py b/exts/StrideSim/StrideSim/base_sample/base_sample.py new file mode 100644 index 0000000..7059751 --- /dev/null +++ b/exts/StrideSim/StrideSim/base_sample/base_sample.py @@ -0,0 +1,127 @@ +# Copyright (c) 2018-2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# +import gc +from abc import abstractmethod + +from omni.isaac.core import World +from omni.isaac.core.scenes.scene import Scene +from omni.isaac.core.utils.stage import create_new_stage_async, update_stage_async + + +class BaseSample: + def __init__(self) -> None: + self._world = None + self._current_tasks = None + self._world_settings = {"physics_dt": 1.0 / 60.0, "stage_units_in_meters": 1.0, "rendering_dt": 1.0 / 60.0} + # self._logging_info = "" + return + + def get_world(self): + return self._world + + def set_world_settings(self, physics_dt=None, stage_units_in_meters=None, rendering_dt=None): + if physics_dt is not None: + self._world_settings["physics_dt"] = physics_dt + if stage_units_in_meters is not None: + self._world_settings["stage_units_in_meters"] = stage_units_in_meters + if rendering_dt is not None: + self._world_settings["rendering_dt"] = rendering_dt + return + + async def load_world_async(self): + """Function called when clicking load button""" + if World.instance() is None: + await create_new_stage_async() + self._world = World(**self._world_settings) + await self._world.initialize_simulation_context_async() + self.setup_scene() + else: + self._world = World.instance() + self._current_tasks = self._world.get_current_tasks() + await self._world.reset_async() + await self._world.pause_async() + await self.setup_post_load() + if len(self._current_tasks) > 0: + self._world.add_physics_callback("tasks_step", self._world.step_async) + return + + async def reset_async(self): + """Function called when clicking reset button""" + if self._world.is_tasks_scene_built() and len(self._current_tasks) > 0: + self._world.remove_physics_callback("tasks_step") + await self._world.play_async() + await update_stage_async() + await self.setup_pre_reset() + await self._world.reset_async() + await self._world.pause_async() + await self.setup_post_reset() + if self._world.is_tasks_scene_built() and len(self._current_tasks) > 0: + self._world.add_physics_callback("tasks_step", self._world.step_async) + return + + @abstractmethod + def setup_scene(self, scene: Scene) -> None: + """used to setup anything in the world, adding tasks happen here for instance. + + Args: + scene (Scene): [description] + """ + return + + @abstractmethod + async def setup_post_load(self): + """called after first reset of the world when pressing load, + initializing private variables happen here. + """ + return + + @abstractmethod + async def setup_pre_reset(self): + """called in reset button before resetting the world + to remove a physics callback for instance or a controller reset + """ + return + + @abstractmethod + async def setup_post_reset(self): + """called in reset button after resetting the world which includes one step with rendering""" + return + + @abstractmethod + async def setup_post_clear(self): + """called after clicking clear button + or after creating a new stage and clearing the instance of the world with its callbacks + """ + return + + # def log_info(self, info): + # self._logging_info += str(info) + "\n" + # return + + def _world_cleanup(self): + self._world.stop() + self._world.clear_all_callbacks() + self._current_tasks = None + self.world_cleanup() + return + + def world_cleanup(self): + """Function called when extension shutdowns and starts again, (hot reloading feature)""" + return + + async def clear_async(self): + """Function called when clicking clear button""" + await create_new_stage_async() + if self._world is not None: + self._world_cleanup() + self._world.clear_instance() + self._world = None + gc.collect() + await self.setup_post_clear() + return diff --git a/exts/StrideSim/StrideSim/base_sample/base_sample_extension.py b/exts/StrideSim/StrideSim/base_sample/base_sample_extension.py new file mode 100644 index 0000000..246dddf --- /dev/null +++ b/exts/StrideSim/StrideSim/base_sample/base_sample_extension.py @@ -0,0 +1,237 @@ +# Copyright (c) 2018-2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# + +import asyncio +import weakref +from abc import abstractmethod + +import omni.ext +import omni.ui as ui +from omni.isaac.core import World +from omni.isaac.examples.base_sample import BaseSample +from omni.isaac.ui.menu import make_menu_item_description +from omni.isaac.ui.ui_utils import btn_builder, get_style, setup_ui_headers # scrolling_frame_builder +from omni.kit.menu.utils import MenuItemDescription, add_menu_items, remove_menu_items + + +class BaseSampleExtension(omni.ext.IExt): + def on_startup(self, ext_id: str): + self._menu_items = None + self._buttons = None + self._ext_id = ext_id + self._sample = None + self._extra_frames = [] + return + + def start_extension( + self, + menu_name: str, + submenu_name: str, + name: str, + title: str, + doc_link: str, + overview: str, + file_path: str, + sample=None, + number_of_extra_frames=1, + window_width=350, + keep_window_open=False, + ): + if sample is None: + self._sample = BaseSample() + else: + self._sample = sample + + menu_items = [make_menu_item_description(self._ext_id, name, lambda a=weakref.proxy(self): a._menu_callback())] + if menu_name == "" or menu_name is None: + self._menu_items = menu_items + elif submenu_name == "" or submenu_name is None: + self._menu_items = [MenuItemDescription(name=menu_name, sub_menu=menu_items)] + else: + self._menu_items = [ + MenuItemDescription( + name=menu_name, sub_menu=[MenuItemDescription(name=submenu_name, sub_menu=menu_items)] + ) + ] + add_menu_items(self._menu_items, "Isaac Examples") + + self._buttons = dict() + self._build_ui( + name=name, + title=title, + doc_link=doc_link, + overview=overview, + file_path=file_path, + number_of_extra_frames=number_of_extra_frames, + window_width=window_width, + keep_window_open=keep_window_open, + ) + return + + @property + def sample(self): + return self._sample + + def get_frame(self, index): + if index >= len(self._extra_frames): + raise Exception(f"there were {len(self._extra_frames)} extra frames created only") + return self._extra_frames[index] + + def get_world(self): + return World.instance() + + def get_buttons(self): + return self._buttons + + def _build_ui( + self, name, title, doc_link, overview, file_path, number_of_extra_frames, window_width, keep_window_open + ): + self._window = omni.ui.Window( + name, width=window_width, height=0, visible=keep_window_open, dockPreference=ui.DockPreference.LEFT_BOTTOM + ) + with self._window.frame: + self._main_stack = ui.VStack(spacing=5, height=0) + with self._main_stack: + setup_ui_headers(self._ext_id, file_path, title, doc_link, overview) + self._controls_frame = ui.CollapsableFrame( + title="World Controls", + width=ui.Fraction(1), + height=0, + collapsed=False, + style=get_style(), + horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED, + vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_ON, + ) + with ui.VStack(style=get_style(), spacing=5, height=0): + for i in range(number_of_extra_frames): + self._extra_frames.append( + ui.CollapsableFrame( + title="", + width=ui.Fraction(0.33), + height=0, + visible=False, + collapsed=False, + style=get_style(), + horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED, + vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_ON, + ) + ) + with self._controls_frame: + with ui.VStack(style=get_style(), spacing=5, height=0): + dict = { + "label": "Load World", + "type": "button", + "text": "Load", + "tooltip": "Load World and Task", + "on_clicked_fn": self._on_load_world, + } + self._buttons["Load World"] = btn_builder(**dict) + self._buttons["Load World"].enabled = True + dict = { + "label": "Reset", + "type": "button", + "text": "Reset", + "tooltip": "Reset robot and environment", + "on_clicked_fn": self._on_reset, + } + self._buttons["Reset"] = btn_builder(**dict) + self._buttons["Reset"].enabled = False + return + + def _set_button_tooltip(self, button_name, tool_tip): + self._buttons[button_name].set_tooltip(tool_tip) + return + + def _on_load_world(self): + async def _on_load_world_async(): + await self._sample.load_world_async() + await omni.kit.app.get_app().next_update_async() + self._sample._world.add_stage_callback("stage_event_1", self.on_stage_event) + self._enable_all_buttons(True) + self._buttons["Load World"].enabled = False + self.post_load_button_event() + self._sample._world.add_timeline_callback("stop_reset_event", self._reset_on_stop_event) + + asyncio.ensure_future(_on_load_world_async()) + return + + def _on_reset(self): + async def _on_reset_async(): + await self._sample.reset_async() + await omni.kit.app.get_app().next_update_async() + self.post_reset_button_event() + + asyncio.ensure_future(_on_reset_async()) + return + + @abstractmethod + def post_reset_button_event(self): + return + + @abstractmethod + def post_load_button_event(self): + return + + @abstractmethod + def post_clear_button_event(self): + return + + def _enable_all_buttons(self, flag): + for btn_name, btn in self._buttons.items(): + if isinstance(btn, omni.ui._ui.Button): + btn.enabled = flag + return + + def _menu_callback(self): + self._window.visible = not self._window.visible + return + + def _on_window(self, status): + # if status: + return + + def on_shutdown(self): + self._extra_frames = [] + if self._sample._world is not None: + self._sample._world_cleanup() + if self._menu_items is not None: + self._sample_window_cleanup() + if self._buttons is not None: + self._buttons["Load World"].enabled = True + self._enable_all_buttons(False) + self.shutdown_cleanup() + return + + def shutdown_cleanup(self): + return + + def _sample_window_cleanup(self): + remove_menu_items(self._menu_items, "Isaac Examples") + self._window = None + self._menu_items = None + self._buttons = None + return + + def on_stage_event(self, event): + if event.type == int(omni.usd.StageEventType.CLOSED): + if World.instance() is not None: + self.sample._world_cleanup() + self.sample._world.clear_instance() + if hasattr(self, "_buttons"): + if self._buttons is not None: + self._enable_all_buttons(False) + self._buttons["Load World"].enabled = True + return + + def _reset_on_stop_event(self, e): + if e.type == int(omni.timeline.TimelineEventType.STOP): + self._buttons["Load World"].enabled = False + self._buttons["Reset"].enabled = True + self.post_clear_button_event() + return diff --git a/exts/StrideSim/StrideSim/quadruped.py b/exts/StrideSim/StrideSim/quadruped.py new file mode 100644 index 0000000..c77af12 --- /dev/null +++ b/exts/StrideSim/StrideSim/quadruped.py @@ -0,0 +1,140 @@ +# Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# + +import numpy as np + +import carb +import omni +import omni.appwindow # Contains handle to keyboard +from omni.isaac.quadruped.robots import Unitree + +from .base_sample import BaseSample + + +class QuadrupedExample(BaseSample): + def __init__(self) -> None: + super().__init__() + self._world_settings["stage_units_in_meters"] = 1.0 + self._world_settings["physics_dt"] = 1.0 / 400.0 + self._world_settings["rendering_dt"] = 5.0 / 400.0 + self._enter_toggled = 0 + self._base_command = [0.0, 0.0, 0.0, 0] + self._event_flag = False + + # bindings for keyboard to command + self._input_keyboard_mapping = { + # forward command + "NUMPAD_8": [5, 0.0, 0.0], + "UP": [5, 0.0, 0.0], + # back command + "NUMPAD_2": [-5, 0.0, 0.0], + "DOWN": [-5, 0.0, 0.0], + # left command + "NUMPAD_6": [0.0, -4.0, 0.0], + "RIGHT": [0.0, -4.0, 0.0], + # right command + "NUMPAD_4": [0.0, 4.0, 0.0], + "LEFT": [0.0, 4.0, 0.0], + # yaw command (positive) + "NUMPAD_7": [0.0, 0.0, 1.0], + "N": [0.0, 0.0, 1.0], + # yaw command (negative) + "NUMPAD_9": [0.0, 0.0, -1.0], + "M": [0.0, 0.0, -1.0], + } + return + + def setup_scene(self) -> None: + world = self.get_world() + self._world.scene.add_default_ground_plane( + z_position=0, + name="default_ground_plane", + prim_path="/World/defaultGroundPlane", + static_friction=0.2, + dynamic_friction=0.2, + restitution=0.01, + ) + self._a1 = world.scene.add( + Unitree( + prim_path="/World/A1", + name="A1", + position=np.array([0, 0, 0.400]), + physics_dt=self._world_settings["physics_dt"], + ) + ) + timeline = omni.timeline.get_timeline_interface() + self._event_timer_callback = timeline.get_timeline_event_stream().create_subscription_to_pop_by_type( + int(omni.timeline.TimelineEventType.STOP), self._timeline_timer_callback_fn + ) + return + + async def setup_post_load(self) -> None: + self._world = self.get_world() + self._appwindow = omni.appwindow.get_default_app_window() + self._input = carb.input.acquire_input_interface() + self._keyboard = self._appwindow.get_keyboard() + self._sub_keyboard = self._input.subscribe_to_keyboard_events(self._keyboard, self._sub_keyboard_event) + self._world.add_physics_callback("sending_actions", callback_fn=self.on_physics_step) + await self._world.play_async() + return + + async def setup_post_reset(self) -> None: + self._event_flag = False + await self._world.play_async() + self._a1.set_state(self._a1._default_a1_state) + self._a1.post_reset() + return + + def on_physics_step(self, step_size) -> None: + if self._event_flag: + self._a1._qp_controller.switch_mode() + self._event_flag = False + self._a1.advance(step_size, self._base_command) + + def _sub_keyboard_event(self, event, *args, **kwargs) -> bool: + """Subscriber callback to when kit is updated.""" + # reset event + self._event_flag = False + # when a key is pressedor released the command is adjusted w.r.t the key-mapping + if event.type == carb.input.KeyboardEventType.KEY_PRESS: + # on pressing, the command is incremented + if event.input.name in self._input_keyboard_mapping: + self._base_command[0:3] += np.array(self._input_keyboard_mapping[event.input.name]) + self._event_flag = True + + # enter, toggle the last command + if event.input.name == "ENTER" and self._enter_toggled is False: + self._enter_toggled = True + if self._base_command[3] == 0: + self._base_command[3] = 1 + else: + self._base_command[3] = 0 + self._event_flag = True + + elif event.type == carb.input.KeyboardEventType.KEY_RELEASE: + # on release, the command is decremented + if event.input.name in self._input_keyboard_mapping: + self._base_command[0:3] -= np.array(self._input_keyboard_mapping[event.input.name]) + self._event_flag = True + # enter, toggle the last command + if event.input.name == "ENTER": + self._enter_toggled = False + # since no error, we are fine :) + return True + + def _timeline_timer_callback_fn(self, event) -> None: + if self._a1: + self._a1.post_reset() + return + + def world_cleanup(self): + self._event_timer_callback = None + if self._world.physics_callback_exists("sending_actions"): + self._world.remove_physics_callback("sending_actions") + return diff --git a/exts/StrideSim/StrideSim/quadruped_extension.py b/exts/StrideSim/StrideSim/quadruped_extension.py new file mode 100644 index 0000000..e789fe7 --- /dev/null +++ b/exts/StrideSim/StrideSim/quadruped_extension.py @@ -0,0 +1,85 @@ +import os + +from .base_sample import BaseSampleExtension +from .quadruped import QuadrupedExample + + +class StrideSimExtension(BaseSampleExtension): + def on_startup(self, ext_id: str): + super().on_startup(ext_id) + + overview = "This Example shows quadruped simulation in Isaac Sim. Currently there is a performance issue with " + overview += ( + "the quadruped gait controller; it's being investigated and will be improved in an upcoming release." + ) + overview += "\n\tKeybord Input:" + overview += "\n\t\tup arrow / numpad 8: Move Forward" + overview += "\n\t\tdown arrow/ numpad 2: Move Reverse" + overview += "\n\t\tleft arrow/ numpad 4: Move Left" + overview += "\n\t\tright arrow / numpad 6: Move Right" + overview += "\n\t\tN / numpad 7: Spin Counterclockwise" + overview += "\n\t\tM / numpad 9: Spin Clockwise" + + overview += "\n\nPress the 'Open in IDE' button to view the source code." + + super().start_extension( + menu_name="", + submenu_name="", + name="StrideSim", + title="Auturbo quadruped robot simulation", + doc_link="https://github.com/AuTURBO/StrideSim", + overview=overview, + file_path=os.path.abspath(__file__), + sample=QuadrupedExample(), + ) + return + + +# import subprocess + +# import omni.ext + + +# # Functions and vars are available to other extension as usual in python: `example.python_ext.some_public_function(x)` +# def some_public_function(x: int): +# print("[StrideSim] some_public_function was called with x: ", x) +# return x**x + + +# # Any class derived from `omni.ext.IExt` in top level module (defined in `python.modules` of `extension.toml`) will be +# # instantiated when extension gets enabled and `on_startup(ext_id)` will be called. Later when extension gets disabled +# # on_shutdown() is called. +# class ExampleExtension(omni.ext.IExt): +# # ext_id is current extension id. It can be used with extension manager to query additional information, like where +# # this extension is located on filesystem. +# def on_startup(self, ext_id): +# print("[StrideSim] startup") + +# self._window = omni.ui.Window("My Window", width=300, height=300) +# with self._window.frame: +# with omni.ui.VStack(): +# label = omni.ui.Label("") + +# def on_train(): +# label.text = "start training" + +# # Execute the specified Python command +# subprocess.run( +# [ +# "python", +# "scripts/rsl_rl/train.py", +# "--task", +# "Isaac-Velocity-Rough-Anymal-D-v0", +# "--headless", +# ] +# ) + +# def on_play(): +# label.text = "empty" + +# with omni.ui.HStack(): +# omni.ui.Button("Train", clicked_fn=on_train) +# omni.ui.Button("Reset", clicked_fn=on_play) + +# def on_shutdown(self): +# print("[StrideSim] shutdown") diff --git a/exts/StrideSim/StrideSim/ui.py b/exts/StrideSim/StrideSim/ui.py deleted file mode 100644 index df415c8..0000000 --- a/exts/StrideSim/StrideSim/ui.py +++ /dev/null @@ -1,48 +0,0 @@ -import subprocess - -import omni.ext - - -# Functions and vars are available to other extension as usual in python: `example.python_ext.some_public_function(x)` -def some_public_function(x: int): - print("[StrideSim] some_public_function was called with x: ", x) - return x**x - - -# Any class derived from `omni.ext.IExt` in top level module (defined in `python.modules` of `extension.toml`) will be -# instantiated when extension gets enabled and `on_startup(ext_id)` will be called. Later when extension gets disabled -# on_shutdown() is called. -class ExampleExtension(omni.ext.IExt): - # ext_id is current extension id. It can be used with extension manager to query additional information, like where - # this extension is located on filesystem. - def on_startup(self, ext_id): - print("[StrideSim] startup") - - self._window = omni.ui.Window("My Window", width=300, height=300) - with self._window.frame: - with omni.ui.VStack(): - label = omni.ui.Label("") - - def on_train(): - label.text = "start training" - - # Execute the specified Python command - subprocess.run( - [ - "python", - "scripts/rsl_rl/train.py", - "--task", - "Isaac-Velocity-Rough-Anymal-D-v0", - "--headless", - ] - ) - - def on_play(): - label.text = "empty" - - with omni.ui.HStack(): - omni.ui.Button("Train", clicked_fn=on_train) - omni.ui.Button("Reset", clicked_fn=on_play) - - def on_shutdown(self): - print("[StrideSim] shutdown")