Merge pull request #268 from maxspahn/ft-reset-mujoco-env

ft[mujoco]: Adds option to randomize obstacles and goals on reset.

examples/mujoco_example.py

@@ -107,7 +107,7 @@ def run_generic_mujoco(
     action_mag = np.random.rand(env.nu) * 1.0
     if render == 'rgb_array':
         env = RecordVideo(env, video_folder=f'{ROBOTMODEL}.mp4')
-    ob, info = env.reset()
+    ob, info = env.reset(options={'randomize_obstacles': False})
 
     t = 0.0
     history = []

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "urdfenvs"
-version = "0.9.12"
+version = "0.9.13"
 description = "Simple simulation environment for robots, based on the urdf files."
 authors = ["Max Spahn <m.spahn@tudelft.nl>"]
 maintainers = [

tests/test_reset.py

@@ -0,0 +1,83 @@
+import numpy as np
+import gymnasium as gym
+
+from robotmodels.utils.robotmodel import RobotModel
+from urdfenvs.sensors.full_sensor import FullSensor
+from mpscenes.obstacles.sphere_obstacle import SphereObstacle
+
+from urdfenvs.scene_examples.goal import goal1
+from urdfenvs.generic_mujoco.generic_mujoco_robot import GenericMujocoRobot
+
+
+obstacle_configuration = {
+    "type": "sphere",
+    "geometry": {"position": [2.0, 2.0, 1.0], "radius": 0.3},
+    "rgba": [0.3, 0.5, 0.6, 1.0],
+    'low': {
+        'position' : [2.0, 2.0, 0.1],
+         'radius': 0.1,
+    },
+    'high': {
+         'position' : [5.0, 5.0, 0.3],
+         'radius': 0.3,
+    },
+}
+obstacle = SphereObstacle(name="sphere_1", content_dict=obstacle_configuration)
+
+def test_reset_mujoco():
+    """
+    Test if the obstacle position changes between two runs of the same environment.
+    """
+    options = {
+        "randomize_obstacles": True,
+        "randomize_goals": True,
+    }
+    sensor = FullSensor(
+        goal_mask=['position'],
+        obstacle_mask=["position", "size"],
+        variance=0.0,
+        physics_engine_name="mujoco",
+    )
+    robot_model = RobotModel("pointRobot", "pointRobot")
+
+    xml_file = robot_model.get_xml_path()
+    robots = [
+        GenericMujocoRobot(xml_file=xml_file, mode="vel"),
+    ]
+    env = gym.make(
+        "generic-mujoco-env-v0",
+        robots=robots,
+        obstacles=[obstacle],
+        goals=[goal1],
+        sensors=[sensor],
+        render=False,
+        enforce_real_time=False,
+    ).unwrapped
+    env.reset()
+    action = np.random.random(env.nu)
+    number_of_steps = 10
+
+    obstacle_position_run_0 = np.zeros(3)
+    obstacle_position_run_1 = np.zeros(3)
+    goal_position_run_0 = np.zeros(3)
+    goal_position_run_1 = np.zeros(3)
+
+    for _ in range(number_of_steps):
+        ob, *_ = env.step(action)
+        observation_run_0 = ob['robot_0']['joint_state']['position']
+        obstacle_position_run_0 = ob['robot_0']['FullSensor']['obstacles'][0]['position']
+        goal_position_run_0 = ob['robot_0']['FullSensor']['goals'][0]['position']
+
+    assert np.allclose(obstacle_position_run_0, obstacle.position())
+    assert np.allclose(goal_position_run_0, goal1.position())
+
+    env.reset(options=options)
+    for _ in range(number_of_steps):
+        ob, *_ = env.step(action)
+        observation_run_1 = ob['robot_0']['joint_state']['position']
+        obstacle_position_run_1 = ob['robot_0']['FullSensor']['obstacles'][0]['position']
+        goal_position_run_1 = ob['robot_0']['FullSensor']['goals'][0]['position']
+
+    assert np.allclose(observation_run_0, observation_run_1)
+    assert not np.allclose(obstacle_position_run_0, obstacle_position_run_1)
+    assert not np.allclose(goal_position_run_0, goal_position_run_1)

urdfenvs/generic_mujoco/generic_mujoco_env.py

@@ -201,6 +201,10 @@ def update_obstacles_position(self):
         for i, obstacle in enumerate(self._obstacles):
             self.data.mocap_pos[i] = obstacle.position(t=self.t)
 
+    def update_goals_position(self):
+        for i, goal in enumerate(self._goals):
+            self.data.site_xpos[i] = goal.position(t=self.t)
+
     def step(self, action: np.ndarray):
         step_start = time.perf_counter()
         self._t += self.dt
@@ -222,6 +226,7 @@ def step(self, action: np.ndarray):
             info = {"Collision": message}
             self._done = True
         self.update_obstacles_position()
+        self.update_goals_position()
         if self.render_mode == "human":
             self.render()
 
@@ -257,6 +262,12 @@ def reset(
         vel: Optional[np.ndarray] = None,
     ):
         super().reset(seed=seed, options=options)
+        if options and options.get("randomize_obstacles", False):
+            for obstacle in self._obstacles:
+                obstacle.shuffle()
+        if options and options.get("randomize_goals", False):
+            for goal in self._goals:
+                goal.shuffle()
         if pos is not None:
             qpos = pos
         else:

urdfenvs/scene_examples/obstacles.py

@@ -10,6 +10,14 @@
     "type": "sphere",
     "geometry": {"position": [2.0, 2.0, 1.0], "radius": 1.0},
     "rgba": [0.3, 0.5, 0.6, 1.0],
+    'low': {
+        'position' : [1.5, 0.0, 0.5],
+        'radius': 0.2,
+    },
+    'high': {
+         'position' : [5.0, 5.0, 1.0],
+        'radius': 0.8,
+    },
 }
 sphereObst1 = SphereObstacle(name="sphere_1", content_dict=obst1Dict)
 obst2Dict = {