isaac sim 模拟仿真代码

主目录

/isaac-sim/extension_examples/user_examples/

__init__.py

from omni.isaac.examples.user_examples.object_control import ObjectControl
from omni.isaac.examples.user_examples.object_control_extension import ObjectControlExtension

object_control_extension.py

import os

from omni.isaac.examples.base_sample import BaseSampleExtension
from omni.isaac.examples.user_examples import ObjectControl


class ObjectControlExtension(BaseSampleExtension):
    def on_startup(self, ext_id: str):
        super().on_startup(ext_id)
        overview = "cube move to [1, 3, 0]"

        super().start_extension(
            menu_name="",
            submenu_name="",
            name="adlong demo",
            title="adlong demo",
            doc_link="https://docs.omniverse.nvidia.com/isaacsim/latest/gui_tutorials/tutorial_advanced_input_devices.html",
            overview=overview,
            file_path=os.path.abspath(__file__),
            sample=ObjectControl(),
        )
        return

object_control.py

from omni.isaac.examples.base_sample import BaseSample
import numpy as np
from omni.isaac.core.objects import DynamicCuboid
from omni.isaac.core.tasks import BaseTask
from omni.isaac.core.controllers import BaseController


class MovingController(BaseController):
    def __init__(self):
        super().__init__(name="moving_controller")
        return

    def forward(self, obj, goal_position):
        pos,ori= obj.get_world_pose()
        size = len(pos)
        for i in range(size):
            if pos[i] < goal_position[i]:
                pos[i] += 0.01
        obj.set_world_pose(pos, ori)
        return


class MovingTask(BaseTask):
    def __init__(self, name, obj, targetPos):
        super().__init__(name=name, offset=None)
        self._obj = obj
        self._goal_position = np.array(targetPos)
        self._task_achieved = False
        return

    def set_up_scene(self, scene):
        super().set_up_scene(scene)
        return
    
    def pre_step(self, control_index, simulation_time):
        cur_position, _ = self._obj.get_world_pose()
        if not self._task_achieved and np.mean(np.abs(self._goal_position - cur_position)) < 0.1:
            self._task_achieved = True
        return
    
    def post_reset(self):
        self._task_achieved = False
        return
    
    def is_done(self):
        return self._task_achieved
    
    def get_goal_position(self):
        return self._goal_position


class ObjectControl(BaseSample):
    def __init__(self) -> None:
        super().__init__()
        self._obj = None
        self._controller = None

    def setup_scene(self):
        world = self.get_world()
        world.scene.add_default_ground_plane()
        temp_usd = DynamicCuboid(
                prim_path="/World/random_cube",
                name="fancy_cube",
                size = 1,
                position=np.array([0.25, 0.25, 0.25]),
                scale=np.array([0.5, 0.5, 0.5]),
                color=np.array([0, 0, 1.0]),
            )
        self._obj = world.scene.add(temp_usd)
        return

    async def setup_post_load(self):
        self._controller = MovingController()
        world = self.get_world()        
        world.add_task(MovingTask(name="my_first_task", obj=self._obj, targetPos=[1, 3, 0]))
        world.add_physics_callback("sim_step", callback_fn=self.physics_step)
        await world.play_async()
        return
    
    async def setup_post_reset(self):
        self._controller.reset()
        world = self.get_world()
        await world.play_async()
        return

    def physics_step(self, step_size):
        world = self.get_world()
        task = world.get_task(name="my_first_task")
        self._controller.forward(self._obj, task.get_goal_position())
        if task.is_done():
            world.pause()
        return

    def world_cleanup(self):
        pass