Skip to content

Commit

Permalink
fix: correct missing observation in 'image' observation mode
Browse files Browse the repository at this point in the history
  • Loading branch information
lilkm committed Nov 24, 2024
1 parent 8cfd64c commit f800cc5
Show file tree
Hide file tree
Showing 5 changed files with 32 additions and 11 deletions.
6 changes: 4 additions & 2 deletions gym_lowcostrobot/envs/lift_cube_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -327,9 +327,11 @@ def step(self, action):
observation = self.get_observation()

# Get the position of the cube and the distance between the end effector and the cube
cube_id = self.model.body("cube").id
cube_pos = self.data.body(cube_id).xpos.copy()
ee_id = self.model.site("end_effector_site").id
ee_pos = self.data.site(ee_id).xpos.copy()
cube_z = observation["cube_pos"][2]
cube_z = cube_pos[2]

# info = {"is_success": self.is_success(ee_pos, observation["cube_pos"])}
info = {}
Expand All @@ -339,7 +341,7 @@ def step(self, action):

# Compute the reward (dense reward)
reward_height = cube_z - self.height_threshold
reward_distance = np.linalg.norm(ee_pos - observation["cube_pos"])
reward_distance = np.linalg.norm(ee_pos - cube_pos)
reward = reward_height + reward_distance
return observation, reward, terminated, truncated, info

Expand Down
8 changes: 6 additions & 2 deletions gym_lowcostrobot/envs/pick_place_cube_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -342,11 +342,15 @@ def step(self, action):
# Get the new observation
observation = self.get_observation()

info = {"is_success": self.is_success(observation["cube_pos"], observation["target_pos"])}
# Get the position of the cube
cube_id = self.model.body("cube").id
cube_pos = self.data.body(cube_id).xpos.copy()

info = {"is_success": self.is_success(cube_pos, observation["target_pos"])}

terminated = info["is_success"]
truncated = False
reward = self.compute_reward(observation["cube_pos"], observation["target_pos"])
reward = self.compute_reward(cube_pos, observation["target_pos"])
return observation, reward, terminated, truncated, info

def goal_distance(self, goal_a, goal_b):
Expand Down
8 changes: 6 additions & 2 deletions gym_lowcostrobot/envs/push_cube_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -334,11 +334,15 @@ def step(self, action):
# Get the new observation
observation = self.get_observation()

info = {"is_success": self.is_success(observation["cube_pos"], observation["target_pos"])}
# Get the position of the cube
cube_id = self.model.body("cube").id
cube_pos = self.data.body(cube_id).xpos.copy()

info = {"is_success": self.is_success(cube_pos, observation["target_pos"])}

terminated = info["is_success"]
truncated = False
reward = self.compute_reward(observation["cube_pos"], observation["target_pos"])
reward = self.compute_reward(cube_pos, observation["target_pos"])
return observation, reward, terminated, truncated, info

def goal_distance(self, goal_a, goal_b):
Expand Down
8 changes: 6 additions & 2 deletions gym_lowcostrobot/envs/reach_cube_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -317,15 +317,19 @@ def step(self, action):
# Get the new observation
observation = self.get_observation()

# Get the position of the cube
cube_id = self.model.body("cube").id
cube_pos = self.data.body(cube_id).xpos.copy()

# Get the position of the cube and the distance between the end effector and the cube
ee_id = self.model.site("end_effector_site").id
ee_pos = self.data.site(ee_id).xpos.copy()

info = {"is_success": self.is_success(ee_pos, observation["cube_pos"])}
info = {"is_success": self.is_success(ee_pos, cube_pos)}

terminated = info["is_success"]
truncated = False
reward = self.compute_reward(ee_pos, observation["cube_pos"])
reward = self.compute_reward(ee_pos, cube_pos)
return observation, reward, terminated, truncated, info

def goal_distance(self, goal_a, goal_b):
Expand Down
13 changes: 10 additions & 3 deletions gym_lowcostrobot/envs/stack_two_cubes_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -330,14 +330,21 @@ def step(self, action):
# Get the new observation
observation = self.get_observation()

# Get the position of the blue cube and the red cube
cube_blue_id = self.model.body("cube_blue").id
cube_blue_pos = self.data.body(cube_blue_id).xpos.copy()

cube_red_id = self.model.body("cube_red").id
cube_red_pos = self.data.body(cube_red_id).xpos.copy()

# The target position is the position of the red cube plus translated in the z-axis by the height of the red cube
target_pos = observation["cube_red_pos"] + np.array([0.0, 0.0, 0.03])
target_pos = cube_red_pos + np.array([0.0, 0.0, 0.03])

info = {"is_success": self.is_success(observation["cube_blue_pos"], target_pos)}
info = {"is_success": self.is_success(cube_blue_pos, target_pos)}

terminated = info["is_success"]
truncated = False
reward = self.compute_reward(observation["cube_blue_pos"], target_pos)
reward = self.compute_reward(cube_blue_pos, target_pos)
return observation, reward, terminated, truncated, info

def goal_distance(self, goal_a, goal_b):
Expand Down

0 comments on commit f800cc5

Please sign in to comment.