From 3567dde0ec9be3f6e7b9b9485e77dbd003b202cf Mon Sep 17 00:00:00 2001
From: Julien Perez <perezjln@users.noreply.github.com>
Date: Thu, 20 Jun 2024 19:38:31 +0200
Subject: [PATCH] add target market to the push task

---
 gym_lowcostrobot/assets/low_cost_robot_6dof/push_cube.xml | 4 +++-
 gym_lowcostrobot/envs/push_cube_env.py                    | 2 ++
 2 files changed, 5 insertions(+), 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
index 58d1f89..9fa8a48 100644
--- 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
index 861d77b..625bd32 100644
--- 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)