Skip to content
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

Refac/UI #58

Merged
merged 2 commits into from
Nov 6, 2024
Merged
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
121 changes: 121 additions & 0 deletions exts/StrideSim/StrideSim/anymal_d_extension.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
import os
import subprocess

import omni.ui as ui
from omni.isaac.ui.ui_utils import btn_builder, get_style, setup_ui_headers # scrolling_frame_builder

from StrideSim.anymal_d import AnymalD
from StrideSim.base_sample import BaseSampleExtension
from StrideSim.parameters import RL_DIR


class AnyamlDExtension(BaseSampleExtension):
Expand Down Expand Up @@ -33,3 +38,119 @@ def on_startup(self, ext_id: str):
sample=AnymalD(),
)
return

def _build_ui(self, title, doc_link, overview, file_path, number_of_extra_frames):
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=True,
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

# New panel for Reinforcement Learning
self._rl_frame = ui.CollapsableFrame(
title="Reinforcement Learning Panel",
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 self._rl_frame:
with ui.VStack(style=get_style(), spacing=5, height=0):
ui.Label("Task Name:", width=ui.Fraction(0.3))
self._rl_task_name_field = ui.StringField(height=0, width=ui.Fraction(0.7))
self._rl_task_name_field.model.set_value("Template-Isaac-Velocity-Rough-Anymal-D-v0")
ui.Label("Window popup:", width=ui.Fraction(0.3))

self._headless_dropdown = ui.ComboBox(0, "on", "off")

dict_train = {
"label": "Train",
"type": "button",
"text": "TRAIN",
"tooltip": "Start training",
"on_clicked_fn": self._on_train,
}
self._buttons["Train"] = btn_builder(**dict_train)

dict_play = {
"label": "Play",
"type": "button",
"text": "PLAY",
"tooltip": "Start playing",
"on_clicked_fn": self._on_play,
}
self._buttons["Play"] = btn_builder(**dict_play)
return

def _on_train(self):
task_name = self._rl_task_name_field.model.get_value_as_string()
training_window = self._headless_dropdown.model.get_item_value_model().get_value_as_string()

command = f"python {RL_DIR}/train.py --task {task_name}"
if training_window == "1":
command += " --headless"

try:
subprocess.Popen(command, shell=True)
print(f"Started training process: {command}")
except Exception as e:
print(f"Error starting training process: {e}")
return

def _on_play(self):
task_name = self._rl_task_name_field.model.get_value_as_string()

command = f"python {RL_DIR}/play.py --task {task_name}"

command += " --num_envs 10"

try:
subprocess.Popen(command, shell=True)
print(f"Started training process: {command}")
except Exception as e:
print(f"Error starting training process: {e}")
return
137 changes: 9 additions & 128 deletions exts/StrideSim/StrideSim/base_sample/base_sample_extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,19 +8,16 @@
#

import asyncio
import subprocess
import weakref
from abc import abstractmethod

import omni.ext
import omni.ui as ui
from omni.isaac.core import World
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

from StrideSim.base_sample import BaseSample
from StrideSim.parameters import RL_DIR


class BaseSampleExtension(omni.ext.IExt):
Expand Down Expand Up @@ -65,15 +62,19 @@ def start_extension(
add_menu_items(self._menu_items, "Isaac Examples")

self._buttons = dict()

self._window = ui.Window(
name, width=window_width, height=0, visible=keep_window_open, dockPreference=ui.DockPreference.LEFT_BOTTOM
)

self._window.deferred_dock_in("Property", ui.DockPolicy.CURRENT_WINDOW_IS_ACTIVE)

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

Expand All @@ -92,99 +93,8 @@ def get_world(self):
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
)

self._window.deferred_dock_in("Property", ui.DockPolicy.CURRENT_WINDOW_IS_ACTIVE)

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=True,
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

# New panel for Reinforcement Learning
self._rl_frame = ui.CollapsableFrame(
title="Reinforcement Learning Panel",
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 self._rl_frame:
with ui.VStack(style=get_style(), spacing=5, height=0):
ui.Label("Task Name:", width=ui.Fraction(0.3))
self._rl_task_name_field = ui.StringField(height=0, width=ui.Fraction(0.7))
self._rl_task_name_field.model.set_value("Template-Isaac-Velocity-Rough-Anymal-D-v0")
ui.Label("Window popup:", width=ui.Fraction(0.3))

self._headless_dropdown = ui.ComboBox(0, "on", "off")

dict_train = {
"label": "Train",
"type": "button",
"text": "TRAIN",
"tooltip": "Start training",
"on_clicked_fn": self._on_train,
}
self._buttons["Train"] = btn_builder(**dict_train)

dict_play = {
"label": "Play",
"type": "button",
"text": "PLAY",
"tooltip": "Start playing",
"on_clicked_fn": self._on_play,
}
self._buttons["Play"] = btn_builder(**dict_play)
@abstractmethod
def _build_ui(self, title, doc_link, overview, file_path, number_of_extra_frames):
return

def _set_button_tooltip(self, button_name, tool_tip):
Expand Down Expand Up @@ -278,32 +188,3 @@ def _reset_on_stop_event(self, e):
self._buttons["Reset"].enabled = True
self.post_clear_button_event()
return

def _on_train(self):
task_name = self._rl_task_name_field.model.get_value_as_string()
training_window = self._headless_dropdown.model.get_item_value_model().get_value_as_string()

command = f"python {RL_DIR}/train.py --task {task_name}"
if training_window == "1":
command += " --headless"

try:
subprocess.Popen(command, shell=True)
print(f"Started training process: {command}")
except Exception as e:
print(f"Error starting training process: {e}")
return

def _on_play(self):
task_name = self._rl_task_name_field.model.get_value_as_string()

command = f"python {RL_DIR}/play.py --task {task_name}"

command += " --num_envs 10"

try:
subprocess.Popen(command, shell=True)
print(f"Started training process: {command}")
except Exception as e:
print(f"Error starting training process: {e}")
return
Loading