Skip to content

Commit

Permalink
Refac/UI (#58)
Browse files Browse the repository at this point in the history 
* refac: ui part를 상속받아 구현하도록 변경

* code formatting
  • Loading branch information
@mqjinwon
mqjinwon authored Nov 6, 2024
1 parent be2f99f commit 607492b
Show file tree
Hide file tree
Showing 2 changed files with 130 additions and 128 deletions.
121 changes: 121 additions & 0 deletions exts/StrideSim/StrideSim/anymal_d_extension.py
View file
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
View file
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

0 comments on commit 607492b

Please sign in to comment.