add target market to the push task

perezjln · Jun 20, 2024 · 3567dde · 3567dde
1 parent 7ce1ec3
commit 3567dde
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/gym_lowcostrobot/assets/low_cost_robot_6dof/push_cube.xml b/gym_lowcostrobot/assets/low_cost_robot_6dof/push_cube.xml
@@ -41,7 +41,6 @@
 			<geom friction="0.5" condim="3" pos="0 0 0" size="0.015 0.015 0.015" type="box" name="cube" rgba="0.5 0 0 1" priority="1"/>
 		</body>
 
-
 		<camera name="camera_front" pos="0.049 0.888 0.317" xyaxes="-0.998 0.056 -0.000 -0.019 -0.335 0.942"/>
 		<camera name="camera_top" pos="0 0 1" euler="0 0 0" mode="fixed"/>
 		<camera name="camera_vizu" pos="-0.1 0.6 0.3" quat="-0.15 -0.1 0.6 1"/>
@@ -122,6 +121,9 @@
 				</body>
 			</body>
 		</body>
+
+		<site name="cube_target" pos="1 0 0" size="0.01" rgba="1 0 0 1" />
+
 	</worldbody>
 
 	<actuator>

diff --git a/gym_lowcostrobot/envs/push_cube_env.py b/gym_lowcostrobot/envs/push_cube_env.py
@@ -229,6 +229,8 @@ def reset(self, seed=None, options=None):
 
         # Sample the target position
         self.target_pos = self.np_random.uniform(self.target_low, self.target_high).astype(np.float32)
+        target_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, "cube_target")
+        self.model.site_pos[target_geom_id] = self.target_pos
 
         # Step the simulation
         mujoco.mj_forward(self.model, self.data)