diff --git a/examples/start_pose.py b/examples/start_pose.py new file mode 100644 index 0000000..f30b187 --- /dev/null +++ b/examples/start_pose.py @@ -0,0 +1,23 @@ +from stretch_mujoco.stretch_mujoco_simulator import StretchMujocoSimulator +from scipy.spatial.transform import Rotation as R + + + +if __name__ == "__main__": + + translation_m = [0.1, -0.2, 0] + + # Convert euler angles to quaternion + euler_angles_degrees = [0, 0, -90] # Example: roll, pitch, yaw (x, y, z) + rotation_obj = R.from_euler('xyz', euler_angles_degrees, degrees=True) + rotation_quat = rotation_obj.as_quat(scalar_first=True) + + + sim = StretchMujocoSimulator( + start_translation=translation_m, + start_rotation_quat=rotation_quat + ) + + sim.start(headless=False) + + while sim.is_running(): ... \ No newline at end of file diff --git a/stretch_mujoco/mujoco_server.py b/stretch_mujoco/mujoco_server.py index 9d4f882..551c500 100644 --- a/stretch_mujoco/mujoco_server.py +++ b/stretch_mujoco/mujoco_server.py @@ -1,6 +1,7 @@ import contextlib from dataclasses import dataclass from multiprocessing.managers import DictProxy, SyncManager +import os import signal import threading import time @@ -192,20 +193,39 @@ def launch_server( stop_mujoco_process_event: threading.Event, data_proxies: MujocoServerProxies, cameras_to_use: list[StretchCameras], + start_translation: list, + start_rotation_quat: list ): - server = cls(scene_xml_path, model, stop_mujoco_process_event, data_proxies) + server = cls(scene_xml_path, model, stop_mujoco_process_event, data_proxies,start_translation , start_rotation_quat) server.run( show_viewer_ui=show_viewer_ui, camera_hz=camera_hz, cameras_to_use=cameras_to_use, ) + def change_start_pose(self,xml_path: str, translation: list, rotation_quat: list): + """Edit the MjSpec and recompile it before loading the model. Mujoco does not allow us to edit body positions at runtime:""" + spec = mujoco.MjSpec.from_file(xml_path) + + current_file_path = os.path.abspath(__file__) + current_directory = os.path.dirname(current_file_path) + spec.meshdir = f"{current_directory}/models/assets/" + spec.texturedir = spec.meshdir + + spec.find_body("base_link").pos = translation + spec.find_body("base_link").quat = rotation_quat + spec.compile() + return MjModel.from_xml_string(spec.to_xml()) + + def __init__( self, scene_xml_path: str | None, model: MjModel | None, stop_mujoco_process_event: threading.Event, data_proxies: MujocoServerProxies, + start_translation: list, + start_rotation_quat: list ): """ Initialize the Simulator handle with a scene @@ -215,11 +235,12 @@ def __init__( """ if scene_xml_path is None: scene_xml_path = utils.default_scene_xml_path - self.mjmodel = MjModel.from_xml_path(scene_xml_path) - elif model is None: - self.mjmodel = MjModel.from_xml_path(scene_xml_path) - if model is not None: - self.mjmodel = model + + if model is None: + model = self.change_start_pose(translation = start_translation, rotation_quat=start_rotation_quat, xml_path=scene_xml_path) + + self.mjmodel = model + self.mjdata = MjData(self.mjmodel) self._base_in_pos_motion = False diff --git a/stretch_mujoco/mujoco_server_passive.py b/stretch_mujoco/mujoco_server_passive.py index e2cb115..864a626 100644 --- a/stretch_mujoco/mujoco_server_passive.py +++ b/stretch_mujoco/mujoco_server_passive.py @@ -56,7 +56,7 @@ def _run_ui_simulation(self, show_viewer_ui: bool): https://mujoco.readthedocs.io/en/stable/python.html#passive-viewer """ self.viewer = mujoco.viewer.launch_passive( - self.mjmodel, self.mjdata, show_left_ui=show_viewer_ui, show_right_ui=show_viewer_ui + model=self.mjmodel, data=self.mjdata, show_left_ui=show_viewer_ui, show_right_ui=show_viewer_ui ) self.viewer._opt.flags[mujoco._enums.mjtVisFlag.mjVIS_RANGEFINDER] = False # Disables the lidar yellow lines. diff --git a/stretch_mujoco/stretch_mujoco_simulator.py b/stretch_mujoco/stretch_mujoco_simulator.py index 4ccd265..bcdff73 100644 --- a/stretch_mujoco/stretch_mujoco_simulator.py +++ b/stretch_mujoco/stretch_mujoco_simulator.py @@ -50,6 +50,8 @@ def __init__( model: MjModel | None = None, camera_hz: float = 30, cameras_to_use: list[StretchCameras] = [], + start_translation: list = [0,0,0], + start_rotation_quat: list = [1,0,0,0] ) -> None: self.scene_xml_path = scene_xml_path self.model = model @@ -57,6 +59,8 @@ def __init__( self.urdf_model = utils.URDFmodel() self._server_process = None self._cameras_to_use = cameras_to_use + self._start_translation = start_translation + self._start_rotation_quat = start_rotation_quat self.is_stop_called = False @@ -106,6 +110,8 @@ def start( self._stop_mujoco_process_event, self.data_proxies, self._cameras_to_use, + self._start_translation, + self._start_rotation_quat, ), daemon=False, # We're gonna handle terminating this in stop_mujoco_process() ) diff --git a/tests/test_one_process.py b/tests/test_one_process.py index a51f39c..c1aa0a2 100644 --- a/tests/test_one_process.py +++ b/tests/test_one_process.py @@ -1,24 +1,29 @@ -from multiprocessing import Manager import signal import threading +from multiprocessing import Manager + +from stretch_mujoco.mujoco_server import MujocoServer, MujocoServerProxies -from stretch_mujoco.mujoco_server import MujocoServer +def main(): + _manager = Manager() + data_proxies = MujocoServerProxies.default(_manager) -from stretch_mujoco.mujoco_server import MujocoServerProxies + event = threading.Event() + signal.signal(signal.SIGTERM, lambda num, frame: event.set()) + signal.signal(signal.SIGINT, lambda num, frame: event.set()) -_manager = Manager() -data_proxies = MujocoServerProxies.default(_manager) + MujocoServer.launch_server( + scene_xml_path=None, + model=None, + camera_hz=30, + show_viewer_ui=True, + stop_mujoco_process_event=event, + data_proxies=data_proxies, + cameras_to_use=[], + start_translation=[0, 0, 0], + start_rotation_quat=[1, 0, 0, 0], + ) -event = threading.Event() -signal.signal(signal.SIGTERM, lambda num, frame: event.set()) -signal.signal(signal.SIGINT, lambda num, frame: event.set()) -MujocoServer.launch_server( - scene_xml_path=None, - model=None, - camera_hz=30, - show_viewer_ui=True, - stop_mujoco_process_event=event, - data_proxies=data_proxies, - cameras_to_use=[] -) +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/tests/test_one_process_passive.py b/tests/test_one_process_passive.py index 99f1642..0632849 100644 --- a/tests/test_one_process_passive.py +++ b/tests/test_one_process_passive.py @@ -1,23 +1,31 @@ -from multiprocessing import Manager import signal import threading +from multiprocessing import Manager from stretch_mujoco.mujoco_server import MujocoServerProxies from stretch_mujoco.mujoco_server_passive import MujocoServerPassive -_manager = Manager() -data_proxies = MujocoServerProxies.default(_manager) +def main(): + _manager = Manager() + data_proxies = MujocoServerProxies.default(_manager) -event = threading.Event() -signal.signal(signal.SIGTERM, lambda num, frame: event.set()) -signal.signal(signal.SIGINT, lambda num, frame: event.set()) + event = threading.Event() + signal.signal(signal.SIGTERM, lambda num, frame: event.set()) + signal.signal(signal.SIGINT, lambda num, frame: event.set()) -MujocoServerPassive.launch_server( - scene_xml_path=None, - model=None, - camera_hz=30, - show_viewer_ui=True, - stop_mujoco_process_event=event, - data_proxies=data_proxies, - cameras_to_use=[] + MujocoServerPassive.launch_server( + scene_xml_path=None, + model=None, + camera_hz=30, + show_viewer_ui=True, + stop_mujoco_process_event=event, + data_proxies=data_proxies, + cameras_to_use=[], + start_translation=[0, 0, 0], + start_rotation_quat=[1, 0, 0, 0], ) + + + +if __name__ == "__main__": + main() \ No newline at end of file