171 lines
7.8 KiB
Python
171 lines
7.8 KiB
Python
import numpy as np
|
|
|
|
|
|
class BeerPongReward:
|
|
def __init__(self):
|
|
|
|
self.robot_collision_objects = ["wrist_palm_link_convex_geom",
|
|
"wrist_pitch_link_convex_decomposition_p1_geom",
|
|
"wrist_pitch_link_convex_decomposition_p2_geom",
|
|
"wrist_pitch_link_convex_decomposition_p3_geom",
|
|
"wrist_yaw_link_convex_decomposition_p1_geom",
|
|
"wrist_yaw_link_convex_decomposition_p2_geom",
|
|
"forearm_link_convex_decomposition_p1_geom",
|
|
"forearm_link_convex_decomposition_p2_geom",
|
|
"upper_arm_link_convex_decomposition_p1_geom",
|
|
"upper_arm_link_convex_decomposition_p2_geom",
|
|
"shoulder_link_convex_decomposition_p1_geom",
|
|
"shoulder_link_convex_decomposition_p2_geom",
|
|
"shoulder_link_convex_decomposition_p3_geom",
|
|
"base_link_convex_geom", "table_contact_geom"]
|
|
|
|
self.cup_collision_objects = ["cup_geom_table3", "cup_geom_table4", "cup_geom_table5", "cup_geom_table6",
|
|
"cup_geom_table7", "cup_geom_table8", "cup_geom_table9", "cup_geom_table10",
|
|
# "cup_base_table", "cup_base_table_contact",
|
|
"cup_geom_table15",
|
|
"cup_geom_table16",
|
|
"cup_geom_table17", "cup_geom1_table8",
|
|
# "cup_base_table_contact",
|
|
# "cup_base_table"
|
|
]
|
|
|
|
|
|
self.ball_traj = None
|
|
self.dists = None
|
|
self.dists_final = None
|
|
self.costs = None
|
|
self.action_costs = None
|
|
self.angle_rewards = None
|
|
self.cup_angles = None
|
|
self.cup_z_axes = None
|
|
self.collision_penalty = 500
|
|
self.reset(None)
|
|
|
|
def reset(self, context):
|
|
self.ball_traj = []
|
|
self.dists = []
|
|
self.dists_final = []
|
|
self.costs = []
|
|
self.action_costs = []
|
|
self.angle_rewards = []
|
|
self.cup_angles = []
|
|
self.cup_z_axes = []
|
|
self.ball_ground_contact = False
|
|
self.ball_table_contact = False
|
|
self.ball_wall_contact = False
|
|
self.ball_cup_contact = False
|
|
|
|
def compute_reward(self, env, action):
|
|
self.ball_id = env.sim.model._body_name2id["ball"]
|
|
self.ball_collision_id = env.sim.model._geom_name2id["ball_geom"]
|
|
self.goal_id = env.sim.model._site_name2id["cup_goal_table"]
|
|
self.goal_final_id = env.sim.model._site_name2id["cup_goal_final_table"]
|
|
self.cup_collision_ids = [env.sim.model._geom_name2id[name] for name in self.cup_collision_objects]
|
|
self.cup_table_id = env.sim.model._body_name2id["cup_table"]
|
|
self.table_collision_id = env.sim.model._geom_name2id["table_contact_geom"]
|
|
self.wall_collision_id = env.sim.model._geom_name2id["wall"]
|
|
self.cup_table_collision_id = env.sim.model._geom_name2id["cup_base_table_contact"]
|
|
self.init_ball_pos_site_id = env.sim.model._site_name2id["init_ball_pos_site"]
|
|
self.ground_collision_id = env.sim.model._geom_name2id["ground"]
|
|
self.robot_collision_ids = [env.sim.model._geom_name2id[name] for name in self.robot_collision_objects]
|
|
|
|
goal_pos = env.sim.data.site_xpos[self.goal_id]
|
|
ball_pos = env.sim.data.body_xpos[self.ball_id]
|
|
ball_vel = env.sim.data.body_xvelp[self.ball_id]
|
|
goal_final_pos = env.sim.data.site_xpos[self.goal_final_id]
|
|
self.dists.append(np.linalg.norm(goal_pos - ball_pos))
|
|
self.dists_final.append(np.linalg.norm(goal_final_pos - ball_pos))
|
|
|
|
action_cost = np.sum(np.square(action))
|
|
self.action_costs.append(action_cost)
|
|
|
|
ball_table_bounce = self._check_collision_single_objects(env.sim, self.ball_collision_id,
|
|
self.table_collision_id)
|
|
|
|
if ball_table_bounce: # or ball_cup_table_cont or ball_wall_con
|
|
self.ball_table_contact = True
|
|
|
|
ball_cup_cont = self._check_collision_with_set_of_objects(env.sim, self.ball_collision_id,
|
|
self.cup_collision_ids)
|
|
if ball_cup_cont:
|
|
self.ball_cup_contact = True
|
|
|
|
ball_wall_cont = self._check_collision_single_objects(env.sim, self.ball_collision_id, self.wall_collision_id)
|
|
if ball_wall_cont and not self.ball_table_contact:
|
|
self.ball_wall_contact = True
|
|
|
|
ball_ground_contact = self._check_collision_single_objects(env.sim, self.ball_collision_id,
|
|
self.ground_collision_id)
|
|
if ball_ground_contact and not self.ball_table_contact:
|
|
self.ball_ground_contact = True
|
|
|
|
self._is_collided = self._check_collision_with_itself(env.sim, self.robot_collision_ids)
|
|
if env._steps == env.ep_length - 1 or self._is_collided:
|
|
|
|
min_dist = np.min(self.dists)
|
|
|
|
ball_in_cup = self._check_collision_single_objects(env.sim, self.ball_collision_id, self.cup_table_collision_id)
|
|
|
|
cost_offset = 0
|
|
|
|
if self.ball_ground_contact: # or self.ball_wall_contact:
|
|
cost_offset += 2
|
|
|
|
if not self.ball_table_contact:
|
|
cost_offset += 2
|
|
|
|
if not ball_in_cup:
|
|
cost_offset += 2
|
|
cost = cost_offset + min_dist ** 2 + 0.5 * self.dists_final[-1] ** 2 + 1e-4 * action_cost # + min_dist ** 2
|
|
else:
|
|
if self.ball_cup_contact:
|
|
cost_offset += 1
|
|
cost = cost_offset + self.dists_final[-1] ** 2 + 1e-4 * action_cost
|
|
|
|
reward = - 1*cost - self.collision_penalty * int(self._is_collided)
|
|
success = ball_in_cup and not self.ball_ground_contact and not self.ball_wall_contact and not self.ball_cup_contact
|
|
else:
|
|
reward = - 1e-4 * action_cost
|
|
success = False
|
|
|
|
infos = {}
|
|
infos["success"] = success
|
|
infos["is_collided"] = self._is_collided
|
|
infos["ball_pos"] = ball_pos.copy()
|
|
infos["ball_vel"] = ball_vel.copy()
|
|
infos["action_cost"] = 5e-4 * action_cost
|
|
|
|
return reward, infos
|
|
|
|
def _check_collision_single_objects(self, sim, id_1, id_2):
|
|
for coni in range(0, sim.data.ncon):
|
|
con = sim.data.contact[coni]
|
|
|
|
collision = con.geom1 == id_1 and con.geom2 == id_2
|
|
collision_trans = con.geom1 == id_2 and con.geom2 == id_1
|
|
|
|
if collision or collision_trans:
|
|
return True
|
|
return False
|
|
|
|
def _check_collision_with_itself(self, sim, collision_ids):
|
|
col_1, col_2 = False, False
|
|
for j, id in enumerate(collision_ids):
|
|
col_1 = self._check_collision_with_set_of_objects(sim, id, collision_ids[:j])
|
|
if j != len(collision_ids) - 1:
|
|
col_2 = self._check_collision_with_set_of_objects(sim, id, collision_ids[j + 1:])
|
|
else:
|
|
col_2 = False
|
|
collision = True if col_1 or col_2 else False
|
|
return collision
|
|
|
|
def _check_collision_with_set_of_objects(self, sim, id_1, id_list):
|
|
for coni in range(0, sim.data.ncon):
|
|
con = sim.data.contact[coni]
|
|
|
|
collision = con.geom1 in id_list and con.geom2 == id_1
|
|
collision_trans = con.geom1 == id_1 and con.geom2 in id_list
|
|
|
|
if collision or collision_trans:
|
|
return True
|
|
return False |