Made walking controller

exts/stride.simulator/stride/simulator/extension.py

@@ -1,24 +1,20 @@
 # TODO: 나중에 UI코드 부분을 빼서 ui folder에 정리해야 함
 
-# Python garbage collenction and asyncronous API
-from threading import Timer
-
 # Omniverse general API
-import pxr
+import omni
 import omni.ext
 import omni.usd
 import omni.kit.ui
 import omni.kit.app
 from omni import ui
 
-from omni.kit.viewport.utility import get_active_viewport
-
 from stride.simulator.interfaces.stride_sim_interface import StrideInterface
 from stride.simulator.vehicles.quadrupedrobot.anymalc import AnymalC, AnymalCConfig
 from stride.simulator.params import SIMULATION_ENVIRONMENTS
 
 import asyncio
 
+
 # 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("[stride.simulator] some_public_function was called with x: ", x)
@@ -37,55 +33,54 @@ def on_startup(self, ext_id):
 
         self._window = ui.Window("Stride Simulator", width=300, height=300)
 
+        self._window.deferred_dock_in("Property", ui.DockPolicy.CURRENT_WINDOW_IS_ACTIVE)
+
         # Start the extension backend
         self._stride_sim = StrideInterface()
 
         with self._window.frame:
             with ui.VStack():
-                label = ui.Label("")
 
-                def on_click():
-                    # Check if we already have a stage loaded (when using autoload feature, it might not be ready yet)
-                    # This is a limitation of the simulator, and we are doing this to make sure that the
-                    # extension does no crash when using the GUI with autoload feature
-                    # If autoload was not enabled, and we are enabling the extension from the Extension widget, then
-                    # we will always have a state open, and the auxiliary timer will never run
+                def on_world():
 
-                    if omni.usd.get_context().get_stage_state() != omni.usd.StageState.CLOSED:
-                        self._stride_sim.initialize_world()
-                    else:
-                        # We need to create a timer to check until the window is properly open and the stage created.
-                        # This is a limitation of the current Isaac Sim simulator and the way it loads extensions :(
-                        self.autoload_helper()
+                    self._stride_sim.initialize_world()
 
                     label.text = "Initialize world"
 
-                    asyncio.ensure_future(
-                        self._stride_sim.load_environment_async(SIMULATION_ENVIRONMENTS["Default Environment"],
-                                                                force_clear=True))
-                    self._stride_sim.initialize_simulation()
+                def on_environment():
+
+                    self._stride_sim.load_asset(SIMULATION_ENVIRONMENTS["Default Environment"], "/World/layout")
 
+                    label.text = "Load environment"
 
-                def on_spawn():
+                def on_simulation():
 
-                    self._anymal_config = AnymalCConfig()
+                    async def respawn():
 
-                    self._anymal = AnymalC(id=ext_id, init_pos=[0.0, 0.0, 0.5],init_orientation=[0.0, 0.0, 0.0, 1.0],
-                                           config=self._anymal_config)
+                        self._anymal_config = AnymalCConfig()
 
-                    label.text = "Open environment"
+                        self._anymal = AnymalC(id=0,
+                                               init_pos=[0.0, 0.0, 0.7],
+                                               init_orientation=[0.0, 0.0, 0.0, 1.0],
+                                               config=self._anymal_config)
+
+                        self._current_tasks = self._stride_sim.world.get_current_tasks()
+                        await self._stride_sim.world.reset_async()
+                        await self._stride_sim.world.pause_async()
+
+                        if len(self._current_tasks) > 0:
+                            self._stride_sim.world.add_physics_callback("tasks_step", self._world.step_async)
+
+                    asyncio.ensure_future(respawn())
+
+                    label.text = "Load simulation"
 
                 with ui.HStack():
-                    ui.Button("Init", clicked_fn=on_click)
-                    ui.Button("Env", clicked_fn=on_spawn)
-
-    def autoload_helper(self):
-        # Check if we already have a viewport and a camera of interest
-        if get_active_viewport() is not None and isinstance(get_active_viewport().stage) == pxr.Usd.Stage and str(
-                get_active_viewport().stage.GetPrimAtPath("/OmniverseKit_Persp")) != "invalid null prim":
-            self._stride_sim.initialize_world()
-        else:
-            Timer(0.1, self.autoload_helper).start()
+                    ui.Button("Init World", clicked_fn=on_world)
+                    ui.Button("Environment", clicked_fn=on_environment)
+                    ui.Button("Simulation", clicked_fn=on_simulation)
+
+                label = ui.Label("")
 
     def on_shutdown(self):
         print("[stride.simulator] stride simulator shutdown")

exts/stride.simulator/stride/simulator/interfaces/stride_sim_interface.py

@@ -10,6 +10,7 @@
 from omni.isaac.core.world import World
 from omni.isaac.core.utils.stage import clear_stage
 from omni.isaac.core.utils.viewports import set_camera_view
+from omni.isaac.core.utils.stage import create_new_stage_async
 from omni.isaac.core.utils import nucleus
 
 # Stride Simulator internal API
@@ -105,17 +106,15 @@ def initialize_world(self):
         """Method that initializes the world object
         """
 
-        self._world = World(**self._world_settings)
+        async def _on_load_world_async():
+            if World.instance() is None:
+                await create_new_stage_async()
+                self._world = World(**self._world_settings)
+                await self._world.initialize_simulation_context_async()
+            else:
+                self._world = World.instance()
 
-    def initialize_simulation(self):
-        """Method that initializes the simulation context
-        """
-
-        # Check if we already have a physics environment activated. If not, then activate it
-        # and only after spawn the vehicle. This is to avoid trying to spawn a vehicle without a physics
-        # environment setup. This way we can even spawn a vehicle in an empty world and it won't care
-        if hasattr(self._world, "_physics_context") is False:
-            asyncio.ensure_future(self._world.initialize_simulation_context_async())
+        asyncio.ensure_future(_on_load_world_async())
 
     def get_vehicle(self, stage_prefix: str):
         """Method that returns the vehicle object given its 'stage_prefix', i.e.,

exts/stride.simulator/stride/simulator/params.py

@@ -41,11 +41,11 @@
     "Stairs Plane": "Terrains/stairs.usd",
 }
 
-ROBOTS_ENVIRONMNETS = {}
+ROBOTS = {}
 
 # Add the Isaac Sim assets to the list
 for asset, path in NVIDIA_SIMULATION_ROBOTS.items():
-    ROBOTS_ENVIRONMNETS[asset] = NVIDIA_ASSETS_PATH + ISAAC_SIM_ROBOTS + "/" + path
+    ROBOTS[asset] = NVIDIA_ASSETS_PATH + ISAAC_SIM_ROBOTS + "/" + path
 
 SIMULATION_ENVIRONMENTS = {}
 
@@ -54,4 +54,4 @@
     SIMULATION_ENVIRONMENTS[asset] = NVIDIA_ASSETS_PATH + ISAAC_SIM_ENVIRONMENTS + "/" + path
 
 # Define the default settings for the simulation environment
-DEFAULT_WORLD_SETTINGS = {"physics_dt": 1.0 / 250.0, "stage_units_in_meters": 1.0, "rendering_dt": 1.0 / 60.0}
+DEFAULT_WORLD_SETTINGS = {"physics_dt": 1.0 / 200.0, "stage_units_in_meters": 1.0, "rendering_dt": 1.0 / 60.0}

exts/stride.simulator/stride/simulator/vehicles/controllers/anymal_controller.py

@@ -1,4 +1,3 @@
-
 import omni
 from omni.isaac.core.utils.nucleus import get_assets_root_path
 
@@ -11,6 +10,7 @@
 
 # TODO: state를 공유해야한다. 아니면 input으로 받아야함 state는 로봇 Vehicle에서 관리하고 있음 e.g) quadrupedrobot.state
 
+
 class AnyamlController(Controller):
     """
     AnyamlController class - It defines a base interface for creating a AnyamlController
@@ -25,31 +25,37 @@ def __init__(self):
         assets_root_path = get_assets_root_path()
 
         # Policy
-        file_content = omni.client.read_file(assets_root_path + "/Isaac/Samples/Quadruped/Anymal_Policies/policy_1.pt")[
-            2
-        ]
+        file_content = omni.client.read_file(assets_root_path +
+                                             "/Isaac/Samples/Quadruped/Anymal_Policies/policy_1.pt")[2]
         file = io.BytesIO(memoryview(file_content).tobytes())
 
         self._policy = torch.jit.load(file)
-        self._base_vel_lin_scale = 2.0
-        self._base_vel_ang_scale = 0.25
-        self._joint_pos_scale = 1.0
-        self._joint_vel_scale = 0.05
+
         self._action_scale = 0.5
-        self._default_joint_pos = np.array([0.0, 0.4, -0.8, 0.0, -0.4, 0.8, -0.0, 0.4, -0.8, -0.0, -0.4, 0.8])
-        self._previous_action = np.zeros(12)
+        self.base_vel_lin_scale = 2.0
+        self.base_vel_ang_scale = 0.25
+        self.joint_pos_scale = 1.0
+        self.joint_vel_scale = 0.05
+        self.default_joint_pos = np.array([0.0, 0.4, -0.8, 0.0, -0.4, 0.8, -0.0, 0.4, -0.8, -0.0, -0.4, 0.8])
+        self.previous_action = np.zeros(12)
         self._policy_counter = 0
 
         # Actuator network
-        file_content = omni.client.read_file(
-            assets_root_path + "/Isaac/Samples/Quadruped/Anymal_Policies/sea_net_jit2.pt"
-        )[2]
+        file_content = omni.client.read_file(assets_root_path +
+                                             "/Isaac/Samples/Quadruped/Anymal_Policies/sea_net_jit2.pt")[2]
         file = io.BytesIO(memoryview(file_content).tobytes())
         self._actuator_network = LstmSeaNetwork()
-        self._actuator_network.setup(file, self._default_joint_pos)
+        self._actuator_network.setup(file, self.default_joint_pos)
         self._actuator_network.reset()
+        self._state = {}  # FIXME: change this variable to state class..
+
+    def get_state(self, joint_positions, joint_velocities):
+        # FIXME: change this variable to state class..
+        self._state["joint_positions"] = joint_positions
+        self._state["joint_velocities"] = joint_velocities
+        pass
 
-    def advance(self, dt, command):
+    def advance(self, dt, obs, command):
         """[summary]
 
         compute the desired torques and apply them to the articulation
@@ -60,13 +66,11 @@ def advance(self, dt, command):
 
         """
         if self._policy_counter % 4 == 0:
-            obs = self._compute_observation(command)
+            used_obs = obs
             with torch.no_grad():
-                obs = torch.from_numpy(obs).view(1, -1).float()
-                self.action = self._policy(obs).detach().view(-1).numpy()
-            self._previous_action = self.action.copy()
-
-        self._dc_interface.wake_up_articulation(self._handle)
+                used_obs = torch.from_numpy(used_obs).view(1, -1).float()
+                self.action = self._policy(used_obs).detach().view(-1).numpy()
+            self.previous_action = self.action.copy()
 
         # joint_state from the DC interface now has the order of
         # 'FL_hip_joint',   'FR_hip_joint',   'RL_hip_joint',   'RR_hip_joint',
@@ -79,15 +83,13 @@ def advance(self, dt, command):
         # RL_hip_joint RL_thigh_joint RL_calf_joint
         # RR_hip_joint RR_thigh_joint RR_calf_joint
         # Convert DC order to controller order for joint info
-        current_joint_pos = self.get_joint_positions()
-        current_joint_vel = self.get_joint_velocities()
+        current_joint_pos = self._state["joint_positions"]
+        current_joint_vel = self._state["joint_velocities"]
         current_joint_pos = np.array(current_joint_pos.reshape([3, 4]).T.flat)
         current_joint_vel = np.array(current_joint_vel.reshape([3, 4]).T.flat)
-        joint_torques, _ = self._actuator_network.compute_torques(
-            current_joint_pos, current_joint_vel, self._action_scale * self.action)
-
-        # finally convert controller order to DC order for command torque
-        torque_reorder = np.array(joint_torques.reshape([4, 3]).T.flat)
-        self._dc_interface.set_articulation_dof_efforts(self._handle, torque_reorder)
+        joint_torques, _ = self._actuator_network.compute_torques(current_joint_pos, current_joint_vel,
+                                                                  self._action_scale * self.action)
 
         self._policy_counter += 1
+
+        return joint_torques

exts/stride.simulator/stride/simulator/vehicles/controllers/controller.py

@@ -2,6 +2,7 @@
 
 # import numpy as np
 
+
 class Controller:
     """
     Controller.
@@ -12,7 +13,7 @@ class Controller:
     def __init__(self):
         self._torque = None
 
-    def advance(self, state):
+    def advance(self, dt, obs, command):
         """
         Add your algorithm to make torque
 

exts/stride.simulator/stride/simulator/vehicles/quadrupedrobot/anymalc.py

@@ -5,9 +5,11 @@
 
 # Get the location of the asset
 from stride.simulator.backends import LoggerBackend
-from stride.simulator.params import ROBOTS_ENVIRONMNETS
+from stride.simulator.params import ROBOTS
 from stride.simulator.vehicles.sensors.imu import Imu
 
+from stride.simulator.vehicles.controllers.anymal_controller import AnyamlController
+
 
 class AnymalCConfig(QuadrupedRobotConfig):
     """
@@ -22,7 +24,7 @@ def __init__(self):
         self.stage_prefix = "/World/AnymalC"
 
         # The USD file that describes the visual aspect of the vehicle
-        self.usd_file = ROBOTS_ENVIRONMNETS["Anymal C"]
+        self.usd_file = ROBOTS["Anymal C"]
 
         # The default sensors for a Anymal C
         self.sensors = [Imu()]  # pylint: disable=use-list-literal FIXME
@@ -45,11 +47,7 @@ def __init__(self, id: int, init_pos, init_orientation, config=AnymalCConfig()):
 
         super().__init__(config.stage_prefix, config.usd_file, id, init_pos, init_orientation, config=config)
 
-
-        # Add callbacks to the physics engine to update each sensor at every timestep and let the sensor decide
-        # depending on its internal update rate whether to generate new data.
-        # TODO: Uncomment this when the physics engine is implemented.
-        # self._world.add_physics_callback(self._stage_prefix + "/Sensors", self.update_sensors)
+        self.controller = AnyamlController()
 
     def update_sensors(self, dt: float):
         """Callback that is called at every physics steps and will call the sensor.update method to generate new
@@ -63,7 +61,21 @@ def update_sensors(self, dt: float):
 
         # Call the update method for the sensor to update its values internally (if applicable)
         for sensor in self._sensors:
-            sensor_data = sensor.update(self._state, dt)
+            # TODO: sensor update 부분 구현 필요
+            try:
+                sensor_data = sensor.update(self._state, dt)
+            except Exception as e:
+                print(f"Error updating sensor: {e}")
+                continue
 
             if sensor_data is not None:
                 print("TODO: Implement backend code.")
+
+    def initialize(self, physics_sim_view=None) -> None:
+        """[summary]
+
+        initialize the dc interface, set up drive mode
+        """
+        super().initialize(physics_sim_view=physics_sim_view)
+        self.get_articulation_controller().set_effort_modes("force")
+        self.get_articulation_controller().switch_control_mode("effort")