diff --git a/examples/mujoco_example.py b/examples/mujoco_example.py index 607115b..c6a1682 100644 --- a/examples/mujoco_example.py +++ b/examples/mujoco_example.py @@ -10,7 +10,7 @@ from urdfenvs.sensors.full_sensor import FullSensor from urdfenvs.sensors.lidar import Lidar from urdfenvs.sensors.sdf_sensor import SDFSensor -from urdfenvs.scene_examples.obstacles import sphereObst1, sphereObst2, wall_obstacles, cylinder_obstacle, dynamicSphereObst1 +from urdfenvs.scene_examples.obstacles import sphereObst1, sphereObst2, wall_obstacles, cylinder_obstacle, dynamicSphereObst1, movable_obstacle from urdfenvs.scene_examples.goal import goal1 from gymnasium.wrappers import RecordVideo @@ -49,7 +49,7 @@ def run_generic_mujoco( else: goal_list = [] if obstacles: - obstacle_list= [sphereObst1, sphereObst2, cylinder_obstacle, dynamicSphereObst1] + wall_obstacles + obstacle_list= [movable_obstacle, sphereObst1, sphereObst2, cylinder_obstacle, dynamicSphereObst1] + wall_obstacles else: obstacle_list= [] full_sensor = FullSensor(['position'], ['position', 'size', 'type', 'orientation'], variance=0.0, physics_engine_name='mujoco') @@ -96,7 +96,7 @@ def run_generic_mujoco( sensors=[lidar_sensor, full_sensor, free_space_decomp, sdf_sensor] else: sensors=[] - env = GenericMujocoEnv( + env: GenericMujocoEnv = GenericMujocoEnv( robots=robots, obstacles=obstacle_list, goals=goal_list, @@ -106,7 +106,7 @@ def run_generic_mujoco( ).unwrapped action_mag = np.random.rand(env.nu) * 1.0 if render == 'rgb_array': - env = RecordVideo(env, video_folder=f'{ROBOTMODEL}.mp4') + env = RecordVideo(env, video_folder=f'{robot_model}.mp4') ob, info = env.reset(options={'randomize_obstacles': False}) t = 0.0 @@ -114,6 +114,7 @@ def run_generic_mujoco( for i in range(n_steps): t0 = time.perf_counter() action = action_mag * np.cos(i/20) + action[-1] = 0.02 ob, _, terminated, _, info = env.step(action) #print(ob['robot_0']) history.append(ob) @@ -128,7 +129,7 @@ def run_generic_mujoco( if __name__ == "__main__": run_generic_mujoco( robot_name='panda', - robot_model='panda_without_gripper', + robot_model='panda_with_gripper', n_steps=int(1e3), render='human', obstacles=True, diff --git a/poetry.lock b/poetry.lock index 84aebf1..de11683 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand. [[package]] name = "absl-py" @@ -946,14 +946,6 @@ files = [ {file = "mapbox_earcut-1.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:9af9369266bf0ca32f4d401152217c46c699392513f22639c6b1be32bde9c1cc"}, {file = "mapbox_earcut-1.0.1-cp311-cp311-win32.whl", hash = "sha256:ff9a13be4364625697b0e0e04ba6a0f77300148b871bba0a85bfa67e972e85c4"}, {file = "mapbox_earcut-1.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:5e736557539c74fa969e866889c2b0149fc12668f35e3ae33667d837ff2880d3"}, - {file = "mapbox_earcut-1.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:4fe92174410e4120022393013705d77cb856ead5bdf6c81bec614a70df4feb5d"}, - {file = "mapbox_earcut-1.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:082f70a865c6164a60af039aa1c377073901cf1f94fd37b1c5610dfbae2a7369"}, - {file = "mapbox_earcut-1.0.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:43d268ece49d0c9e22cb4f92cd54c2cc64f71bf1c5e10800c189880d923e1292"}, - {file = "mapbox_earcut-1.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7748f1730fd36dd1fcf0809d8f872d7e1ddaa945f66a6a466ad37ef3c552ae93"}, - {file = "mapbox_earcut-1.0.1-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:5a82d10c8dec2a0bd9a6a6c90aca7044017c8dad79f7e209fd0667826f842325"}, - {file = "mapbox_earcut-1.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:01b292588cd3f6bad7d76ee31c004ed1b557a92bbd9602a72d2be15513b755be"}, - {file = "mapbox_earcut-1.0.1-cp312-cp312-win32.whl", hash = "sha256:fce236ddc3a56ea7260acc94601a832c260e6ac5619374bb2cec2e73e7414ff0"}, - {file = "mapbox_earcut-1.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:1ce86407353b4f09f5778c436518bbbc6f258f46c5736446f25074fe3d3a3bd8"}, {file = "mapbox_earcut-1.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:aa6111a18efacb79c081f3d3cdd7d25d0585bb0e9f28896b207ebe1d56efa40e"}, {file = "mapbox_earcut-1.0.1-cp36-cp36m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f2911829d1e6e5e1282fbe2840fadf578f606580f02ed436346c2d51c92f810b"}, {file = "mapbox_earcut-1.0.1-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:01ff909a7b8405a923abedd701b53633c997cc2b5dc9d5b78462f51c25ec2c33"}, @@ -1546,13 +1538,13 @@ validation = ["lxml"] [[package]] name = "pyglet" -version = "1.5.28" +version = "1.5.29" description = "Cross-platform windowing and multimedia library" optional = false python-versions = "*" files = [ - {file = "pyglet-1.5.28-py3-none-any.whl", hash = "sha256:ea312a553c266a0f45a9c209f565e3c02371c83d89fd8931422c6475ce4ecbae"}, - {file = "pyglet-1.5.28.zip", hash = "sha256:68bff532226b0e8f54dfcc2f6df7d0e463b451fae97fe9fa4af4131d34afbc00"}, + {file = "pyglet-1.5.29-py3-none-any.whl", hash = "sha256:c9bf6deab6add99e79d635dd74c2bf5a646f2c7fb1859d71f5042fb481d85cb6"}, + {file = "pyglet-1.5.29.zip", hash = "sha256:397f40ebf147e7f63b1ba808b96daba1fa0c9e3855c98c9662f8c37325ff05b1"}, ] [[package]] @@ -1848,7 +1840,6 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, @@ -1921,13 +1912,13 @@ use-chardet-on-py3 = ["chardet (>=3.0.2,<6)"] [[package]] name = "robotmodels" -version = "0.1.3" +version = "0.1.9" description = "A collection of robot models" optional = false python-versions = "<4.0,>=3.8" files = [ - {file = "robotmodels-0.1.3-py3-none-any.whl", hash = "sha256:b5a25525b0c2fd73a9096c11c511b1ba9f3c4c8b0971ce3e96a6dd20f6c64cdf"}, - {file = "robotmodels-0.1.3.tar.gz", hash = "sha256:7e57de382621b2b608e7d8013d8082fe3eede3c610b0a52b84a1570b9e12a745"}, + {file = "robotmodels-0.1.9-py3-none-any.whl", hash = "sha256:2bd4c6c5af13b277b0e7028ae122a4489df15aff9bb8ad29566fdbff700caf5d"}, + {file = "robotmodels-0.1.9.tar.gz", hash = "sha256:162f68ef4aa8de048c278e8c657718fd5cd057c8eb9501f5c7a94ff74b7c3b00"}, ] [package.dependencies] @@ -2654,4 +2645,4 @@ keyboard = [] [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "772ea61ecf8866a9fdc78927640561a05b95c5d39c522b78c15504be3a5eafcb" +content-hash = "d079bea33e5975ca9ee2abf5e1d9ea131f27f75511d1dc0234bfdb85af279389" diff --git a/pyproject.toml b/pyproject.toml index 3291050..a40cb09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "urdfenvs" -version = "0.9.13" +version = "0.9.14" description = "Simple simulation environment for robots, based on the urdf files." authors = ["Max Spahn "] maintainers = [ @@ -23,10 +23,10 @@ deprecation = "^2.1.0" mpscenes = "^0.4.4" gymnasium = "1.0.0.a1" dill = "^0.3.7" -robotmodels = "^0.1.3" scipy = "^1.9.0" moviepy = "^1.0.3" trimesh = {version = "^3.0,<3.22", extras = ['easy']} +robotmodels = "^0.1.9" [tool.poetry.extras] keyboard = ["pynput", "multiprocess"] diff --git a/tests/test_full_sensor.py b/tests/test_full_sensor.py index 24ab184..0bf4c2f 100644 --- a/tests/test_full_sensor.py +++ b/tests/test_full_sensor.py @@ -8,7 +8,7 @@ from urdfenvs.sensors.full_sensor import FullSensor from urdfenvs.urdf_common.urdf_env import UrdfEnv -from urdfenvs.scene_examples.obstacles import sphereObst1, dynamicSphereObst3 +from urdfenvs.scene_examples.obstacles import sphereObst1, dynamicSphereObst3, movable_obstacle from urdfenvs.scene_examples.goal import goal1 from urdfenvs.robots.generic_urdf import GenericUrdfReacher from urdfenvs.generic_mujoco.generic_mujoco_env import GenericMujocoEnv @@ -24,6 +24,7 @@ def test_full_sensor(): env.reset() env.add_obstacle(sphereObst1) env.add_obstacle(dynamicSphereObst3) + env.add_obstacle(movable_obstacle) sensor = FullSensor( goal_mask=["position", "is_primary_goal"], obstacle_mask=["velocity", "position", "size"], @@ -41,8 +42,10 @@ def test_full_sensor(): assert isinstance(full_sensor_ob["obstacles"], dict) assert isinstance(full_sensor_ob["obstacles"][2], dict) assert isinstance(full_sensor_ob["obstacles"][2]["position"], np.ndarray) + assert isinstance(full_sensor_ob["obstacles"][3], dict) + assert isinstance(full_sensor_ob["obstacles"][3]["position"], np.ndarray) assert isinstance(full_sensor_ob["goals"], dict) - assert isinstance(full_sensor_ob["goals"][4]["position"], np.ndarray) + assert isinstance(full_sensor_ob["goals"][5]["position"], np.ndarray) np.testing.assert_array_almost_equal( full_sensor_ob["obstacles"][2]["velocity"], sphereObst1.velocity(t=env.t()), @@ -54,12 +57,12 @@ def test_full_sensor(): decimal=4, ) np.testing.assert_array_almost_equal( - full_sensor_ob["goals"][4]["position"], + full_sensor_ob["goals"][5]["position"], goal1.position(t=env.t()), decimal=4, ) np.testing.assert_array_almost_equal( - full_sensor_ob["goals"][4]["is_primary_goal"], + full_sensor_ob["goals"][5]["is_primary_goal"], goal1.is_primary_goal(), decimal=4, ) @@ -83,6 +86,28 @@ def test_full_sensor(): dynamicSphereObst3.radius(), decimal=2, ) + # Estimate the position of the movable obstacle after 10 steps + pos = movable_obstacle.position(t=0) + vel = movable_obstacle.velocity(t=0) + acc = np.array([0, 0, -9.81]) + for i in range(10): + vel += acc * env.dt + pos += vel * env.dt + np.testing.assert_array_almost_equal( + full_sensor_ob["obstacles"][4]["position"], + pos, + decimal=1, + ) + np.testing.assert_array_almost_equal( + full_sensor_ob["obstacles"][4]["velocity"], + vel, + decimal=1, + ) + np.testing.assert_array_almost_equal( + full_sensor_ob["obstacles"][4]["size"], + movable_obstacle.size(), + decimal=2, + ) env.close() @@ -93,11 +118,7 @@ def test_full_sensor_mujoco(): variance=0.0, physics_engine_name="mujoco", ) - if os.path.exists("pointRobot"): - shutil.rmtree("pointRobot") - robot_model_original = RobotModel("pointRobot", "pointRobot") - robot_model_original.copy_model(os.path.join(os.getcwd(), "pointRobot")) - robot_model = LocalRobotModel("pointRobot", "pointRobot") + robot_model = RobotModel("pointRobot", "pointRobot") xml_file = robot_model.get_xml_path() robots = [ @@ -106,7 +127,7 @@ def test_full_sensor_mujoco(): env: GenericMujocoEnv = gym.make( "generic-mujoco-env-v0", robots=robots, - obstacles=[sphereObst1, dynamicSphereObst3], + obstacles=[sphereObst1, dynamicSphereObst3, movable_obstacle], goals=[goal1], sensors=[sensor], render=False, @@ -146,11 +167,28 @@ def test_full_sensor_mujoco(): np.testing.assert_array_almost_equal( full_sensor_ob["obstacles"][1]["position"], dynamicSphereObst3.position(t=env.t), - decimal=4, + decimal=2, ) np.testing.assert_array_almost_equal( full_sensor_ob["obstacles"][1]["size"], dynamicSphereObst3.radius(), decimal=2, ) + # Estimate the position of the movable obstacle after 10 steps + pos = movable_obstacle.position(t=0) + vel = movable_obstacle.velocity(t=0) + acc = np.array([0, 0, -9.81]) + for i in range(10): + vel += acc * env.dt + pos += vel * env.dt + np.testing.assert_array_almost_equal( + full_sensor_ob["obstacles"][2]["position"], + pos, + decimal=1, + ) + np.testing.assert_array_almost_equal( + full_sensor_ob["obstacles"][2]["size"], + movable_obstacle.size(), + decimal=2, + ) env.close() diff --git a/urdfenvs/generic_mujoco/generic_mujoco_env.py b/urdfenvs/generic_mujoco/generic_mujoco_env.py index 0126e1d..be58ce7 100644 --- a/urdfenvs/generic_mujoco/generic_mujoco_env.py +++ b/urdfenvs/generic_mujoco/generic_mujoco_env.py @@ -28,6 +28,7 @@ DEFAULT_SIZE = 480 + class LinkIdNotFoundError(Exception): pass @@ -42,6 +43,7 @@ class GenericMujocoEnv(Env): "render_fps": 20, } _t: float + _number_movable_obstacles: int def __init__( self, @@ -67,16 +69,16 @@ def __init__( self._sensors = [] else: self._sensors = sensors + self._number_movable_obstacles = 0 self.add_scene() self._t = 0.0 - render_mode = render if render is True: - render_mode = 'human' + render_mode = "human" if render is False: render_mode = None - if enforce_real_time and render_mode == 'human': + if enforce_real_time and render_mode == "human": self._enforce_real_time = True else: self._enforce_real_time = False @@ -84,7 +86,6 @@ def __init__( self.width = width self.height = height - self.frame_skip = frame_skip self._initialize_simulation() @@ -109,41 +110,49 @@ def __init__( self.camera_id = camera_id self.mujoco_renderer = MujocoRenderer( - self.model, self.data, DEFAULT_CAMERA_CONFIG, + self.model, + self.data, + DEFAULT_CAMERA_CONFIG, height=DEFAULT_SIZE, width=DEFAULT_SIZE, ) def render(self): - return self.mujoco_renderer.render( - self.render_mode - ) + return self.mujoco_renderer.render(self.render_mode) def get_observation_space(self) -> gym.spaces.Dict: """Get observation space.""" - observation_space_as_dict= dict(gym.spaces.Dict({ - "joint_state": gym.spaces.Dict( + observation_space_as_dict = dict( + gym.spaces.Dict( { - "position": gym.spaces.Box( - low=self.joint_limits()[:, 0], - high=self.joint_limits()[:, 1], - dtype=float, - ), - "velocity": gym.spaces.Box( - low=np.ones_like(self.joint_limits()[:, 0]) - * -10.0, - high=np.ones_like(self.joint_limits()[:, 0]) - * 10.0, - dtype=float, - ), + "joint_state": gym.spaces.Dict( + { + "position": gym.spaces.Box( + low=self.joint_limits()[:, 0], + high=self.joint_limits()[:, 1], + dtype=float, + ), + "velocity": gym.spaces.Box( + low=np.ones_like(self.joint_limits()[:, 0]) + * -10.0, + high=np.ones_like(self.joint_limits()[:, 0]) + * 10.0, + dtype=float, + ), + } + ) } ) - })) + ) for sensor in self._sensors: observation_space_as_dict.update( - sensor.get_observation_space(self.obstacles_dict, self.goals_dict) + sensor.get_observation_space( + self.obstacles_dict, self.goals_dict + ) ) - observation_space = gym.spaces.Dict({'robot_0': gym.spaces.Dict(observation_space_as_dict)}) + observation_space = gym.spaces.Dict( + {"robot_0": gym.spaces.Dict(observation_space_as_dict)} + ) return observation_space def get_action_space(self) -> np.ndarray: @@ -154,7 +163,7 @@ def get_action_space(self) -> np.ndarray: ) def joint_limits(self) -> np.ndarray: - return self.model.jnt_range + return self.model.jnt_range[: self.nq] def actuator_limits(self) -> np.ndarray: return self.model.actuator_ctrlrange @@ -166,21 +175,19 @@ def _initialize_simulation( self, ): - file_name = self._xml_file.split('/')[-1] + file_name = self._xml_file.split("/")[-1] mjcf.export_with_assets( self._model_dm, - 'xml_model', + "xml_model", out_file_name=file_name, ) - self.model = mujoco.MjModel.from_xml_path(f'xml_model/{file_name}') - + self.model = mujoco.MjModel.from_xml_path(f"xml_model/{file_name}") self.model.body_pos[0] = [0, 1, 1] self.model.vis.global_.offwidth = self.width self.model.vis.global_.offheight = self.height self.data = mujoco.MjData(self.model) - @property def dt(self) -> float: return self.model.opt.timestep * self.frame_skip @@ -198,7 +205,11 @@ def goals_dict(self) -> dict: return {k: v for k, v in enumerate(self._goals)} def update_obstacles_position(self): - for i, obstacle in enumerate(self._obstacles): + + non_movable_obstacles = [ + obstacle for obstacle in self._obstacles if not obstacle.movable() + ] + for i, obstacle in enumerate(non_movable_obstacles): self.data.mocap_pos[i] = obstacle.position(t=self.t) def update_goals_position(self): @@ -219,10 +230,14 @@ def step(self, action: np.ndarray): for contact in self.data.contact: body1 = self.model.geom(contact.geom1).name body2 = self.model.geom(contact.geom2).name + if "movable" in body1 or "movable" in body2: + continue - message = f"Collision occured at {round(self.t, 2)} " \ - f"between {body1} and obstacle " \ + message = ( + f"Collision occured at {round(self.t, 2)} " + f"between {body1} and obstacle " f"with id {body2}" + ) info = {"Collision": message} self._done = True self.update_obstacles_position() @@ -244,7 +259,7 @@ def step(self, action: np.ndarray): time.sleep(sleep_time) step_final_end = time.perf_counter() total_step_time = step_final_end - step_start - real_time_factor = self.dt/total_step_time + real_time_factor = self.dt / total_step_time logging.info(f"Real time factor {real_time_factor}") return ( ob, @@ -275,15 +290,15 @@ def reset( if vel is not None: qvel = vel else: - qvel = np.zeros(self.model.nv) + qvel = np.zeros(self.nv) self.set_state(qpos, qvel) self._t = 0.0 return self._get_obs(), {} def set_state(self, qpos, qvel): - assert qpos.shape == (self.model.nq,) and qvel.shape == (self.model.nv,) - self.data.qpos[:] = np.copy(qpos) - self.data.qvel[:] = np.copy(qvel) + assert qpos.shape == (self.nq,) and qvel.shape == (self.nv,) + self.data.qpos[: self.nq] = np.copy(qpos) + self.data.qvel[: self.nv] = np.copy(qvel) if self.model.na == 0: self.data.act[:] = None mujoco.mj_forward(self.model, self.data) @@ -294,11 +309,11 @@ def nu(self) -> int: @property def nq(self) -> int: - return self.model.nq + return self.model.nq - 7 * self._number_movable_obstacles @property def nv(self) -> int: - return self.model.nv + return self.model.nv - 6 * self._number_movable_obstacles def do_simulation(self, ctrl, n_frames) -> None: """ @@ -331,10 +346,9 @@ def add_scene(self) -> None: if isinstance(sensor, Lidar): self.add_range_finder(sensor) - def add_range_finder(self, sensor: Lidar) -> None: try: - lidar_body = self._model_dm.find('body', sensor._link_name) + lidar_body = self._model_dm.find("body", sensor._link_name) except Exception as e: print(e) raise LinkIdNotFoundError( @@ -342,40 +356,40 @@ def add_range_finder(self, sensor: Lidar) -> None: ) for i in range(sensor._nb_rays): - angle = i/sensor._nb_rays * np.pi * 2 - start_position = np.array([ - np.cos(angle) * 0.00, - np.sin(angle) * 0.00, - 0.0, - ]) - end_position = np.array([ - np.cos(angle) * 0.01, - np.sin(angle) * 0.01, - 0.0, - ]) - fromto_string = " ".join(map(str, start_position)) + angle = i / sensor._nb_rays * np.pi * 2 + start_position = np.array( + [ + np.cos(angle) * 0.00, + np.sin(angle) * 0.00, + 0.0, + ] + ) + end_position = np.array( + [ + np.cos(angle) * 0.01, + np.sin(angle) * 0.01, + 0.0, + ] + ) + fromto_string = " ".join(map(str, start_position)) fromto_string += " " + " ".join(map(str, end_position)) site_values = { - "name":f"{sensor.name()}_rf_{i}", - "type":"capsule", - "size":"0.01", - "fromto": fromto_string, + "name": f"{sensor.name()}_rf_{i}", + "type": "capsule", + "size": "0.01", + "fromto": fromto_string, } lidar_body.add("site", **site_values) - rangefinder_values = { - 'name': f'{sensor.name()}_{i}', - 'site': f'{sensor.name()}_rf_{i}', - 'cutoff': str(sensor._ray_length + 0.01/2), + "name": f"{sensor.name()}_{i}", + "site": f"{sensor.name()}_rf_{i}", + "cutoff": str(sensor._ray_length + 0.01 / 2), } self._model_dm.sensor.add("rangefinder", **rangefinder_values) - - def add_obstacle( - self, obst: CollisionObstacle - ) -> None: + def add_obstacle(self, obst: CollisionObstacle) -> None: if obst.type() == "sphere": size = str(obst.size()[0]) else: @@ -388,14 +402,16 @@ def add_obstacle( "size": size, } obstacle_body = self._model_dm.worldbody.add( - "body", - name=obst.name(), - pos=obst.position().tolist(), - mocap=True, + "body", + name=obst.name(), + pos=obst.position().tolist(), + mocap=not obst.movable(), ) + if obst.movable(): + obstacle_body.add("freejoint", name=f"freejoint_{obst.name()}") + self._number_movable_obstacles += 1 obstacle_body.add("geom", **geom_values) - def add_sub_goal(self, sub_goal: SubGoal) -> None: position = np.zeros(3) for index in sub_goal.indices(): @@ -412,12 +428,18 @@ def add_sub_goal(self, sub_goal: SubGoal) -> None: def _get_obs(self): observation = { "joint_state": { - "position": self.data.qpos, - "velocity": self.data.qvel, + "position": self.data.qpos[0 : self.nq], + "velocity": self.data.qvel[0 : self.nv], } } for sensor in self._sensors: - observation.update({sensor.name(): sensor.sense(0, self.obstacles_dict, self.goals_dict, self.t)}) + observation.update( + { + sensor.name(): sensor.sense( + 0, self.obstacles_dict, self.goals_dict, self.t + ) + } + ) return {"robot_0": observation} def close(self): @@ -427,8 +449,8 @@ def close(self): def xml(self) -> str: """Return the xml string of the model. - As the model from the xml file may be changed during the simulation, - by adding obstacles, goals, sensors, etc., this method ensures that + As the model from the xml file may be changed during the simulation, + by adding obstacles, goals, sensors, etc., this method ensures that the updated model can be saved as an xml file. """ diff --git a/urdfenvs/sensors/full_sensor.py b/urdfenvs/sensors/full_sensor.py index 9cefac4..1f8d347 100644 --- a/urdfenvs/sensors/full_sensor.py +++ b/urdfenvs/sensors/full_sensor.py @@ -122,13 +122,25 @@ def sense(self, robot, obstacles: dict, goals: dict, t: float): observation = {} for mask_item in self._obstacle_mask: if mask_item == "position": - value, _ = self._physics_engine.get_obstacle_pose(obst_id) + value, _ = self._physics_engine.get_obstacle_pose( + obst_id, + obstacle.name(), + movable=obstacle.movable(), + ) elif mask_item == "orientation": - _, value_raw = self._physics_engine.get_obstacle_pose(obst_id) + _, value_raw = self._physics_engine.get_obstacle_pose( + obst_id, + obstacle.name(), + movable=obstacle.movable(), + ) value = value_raw[3:] + value_raw[:3] elif mask_item == "velocity": - value, _ = self._physics_engine.get_obstacle_velocity(obst_id) + value, _ = self._physics_engine.get_obstacle_velocity( + obst_id, + obstacle.name(), + movable=obstacle.movable(), + ) else: try: diff --git a/urdfenvs/sensors/physics_engine_interface.py b/urdfenvs/sensors/physics_engine_interface.py index 5033b08..62cb9c6 100644 --- a/urdfenvs/sensors/physics_engine_interface.py +++ b/urdfenvs/sensors/physics_engine_interface.py @@ -20,11 +20,11 @@ def __init__(self): pass @abstractmethod - def get_obstacle_pose(self, *args) -> Tuple[List[float], List[float]]: + def get_obstacle_pose(self, *args, **kwargs) -> Tuple[List[float], List[float]]: pass @abstractmethod - def get_obstacle_velocity(self, *args) -> Tuple[List[float], List[float]]: + def get_obstacle_velocity(self, *args, **kwargs) -> Tuple[List[float], List[float]]: pass @abstractmethod @@ -83,11 +83,11 @@ def extract_link_id(self, robot, link_name: str) -> int: ) - def get_obstacle_pose(self, obst_id: int) -> Tuple[List[float], List[float]]: + def get_obstacle_pose(self, obst_id: int, obst_name: str, movable: bool = False) -> None: position, orientation = pybullet.getBasePositionAndOrientation(obst_id) return position, orientation - def get_obstacle_velocity(self, obst_id: int) -> Tuple[List[float], List[float]]: + def get_obstacle_velocity(self, obst_id: int, obst_name: str, movable: bool = False) -> None: linear, angular = pybullet.getBaseVelocity(obst_id) return linear, angular @@ -144,9 +144,14 @@ def get_link_orientation(self, robot, link_id) -> np.ndarray: link_orientation_matrix = np.reshape(self._data.xmat[link_id], (3, 3)) return Rotation.from_matrix(link_orientation_matrix).as_quat() - def get_obstacle_pose(self, obst_id: int) -> Tuple[List[float], List[float]]: - pos = self._data.mocap_pos[obst_id] - ori = self._data.mocap_quat[obst_id] + def get_obstacle_pose(self, obst_id: int, obst_name: str, movable: bool = False) -> Tuple[List[float], List[float]]: + if movable: + + free_joint_data = self._data.jnt(f"freejoint_{obst_name}").qpos + print(free_joint_data) + return free_joint_data[0:3].tolist(), free_joint_data[3:].tolist() + pos = self._data.body(obst_name).xpos + ori = self._data.body(obst_name).xquat return pos.tolist(), ori.tolist() def get_goal_pose(self, goal_id: int) -> Tuple[List[float]]: @@ -171,7 +176,7 @@ def ray_cast( ray_value = self._data.sensordata[ray_index] - (0.01 / 2) return ray_value - def get_obstacle_velocity(self, obst_id: int) -> None: + def get_obstacle_velocity(self, obst_id: int, obst_name: str, movable: bool = False) -> None: raise NotImplementedError( "Obstacle velocity not implemented for mujoco." )