first clean up and some non working ideas sketched
This commit is contained in:
@@ -1,11 +1,11 @@
|
||||
import mujoco_py.builder
|
||||
import os
|
||||
|
||||
import mujoco_py.builder
|
||||
import numpy as np
|
||||
from gym import utils, spaces
|
||||
from gym import utils
|
||||
from gym.envs.mujoco import MujocoEnv
|
||||
from alr_envs.alr.mujoco.beerpong.beerpong_reward_staged import BeerPongReward
|
||||
|
||||
from alr_envs.alr.mujoco.beerpong.beerpong_reward_staged import BeerPongReward
|
||||
|
||||
CUP_POS_MIN = np.array([-1.42, -4.05])
|
||||
CUP_POS_MAX = np.array([1.42, -1.25])
|
||||
@@ -18,12 +18,14 @@ CUP_POS_MAX = np.array([1.42, -1.25])
|
||||
# CUP_POS_MIN = np.array([-0.16, -2.2])
|
||||
# CUP_POS_MAX = np.array([0.16, -1.7])
|
||||
|
||||
|
||||
class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
|
||||
def __init__(self, frame_skip=1, apply_gravity_comp=True, noisy=False,
|
||||
rndm_goal=False, cup_goal_pos=None):
|
||||
def __init__(
|
||||
self, frame_skip=1, apply_gravity_comp=True, noisy=False,
|
||||
rndm_goal=False, cup_goal_pos=None
|
||||
):
|
||||
|
||||
cup_goal_pos = np.array(cup_goal_pos if cup_goal_pos is not None else [-0.3, -1.2, 0.840])
|
||||
if cup_goal_pos.shape[0]==2:
|
||||
if cup_goal_pos.shape[0] == 2:
|
||||
cup_goal_pos = np.insert(cup_goal_pos, 2, 0.840)
|
||||
self.cup_goal_pos = np.array(cup_goal_pos)
|
||||
|
||||
@@ -50,7 +52,7 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
|
||||
# self._release_step = 130 # time step of ball release
|
||||
self.release_step = 100 # time step of ball release
|
||||
|
||||
self.ep_length = 600//frame_skip
|
||||
self.ep_length = 600 // frame_skip
|
||||
self.cup_table_id = 10
|
||||
|
||||
if noisy:
|
||||
@@ -71,14 +73,6 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
|
||||
def start_vel(self):
|
||||
return self._start_vel
|
||||
|
||||
@property
|
||||
def current_pos(self):
|
||||
return self.sim.data.qpos[0:7].copy()
|
||||
|
||||
@property
|
||||
def current_vel(self):
|
||||
return self.sim.data.qvel[0:7].copy()
|
||||
|
||||
def reset(self):
|
||||
self.reward_function.reset(self.add_noise)
|
||||
return super().reset()
|
||||
@@ -122,7 +116,7 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
|
||||
self.sim.data.qpos[7::] = self.sim.data.site_xpos[self.ball_site_id, :].copy()
|
||||
self.sim.data.qvel[7::] = self.sim.data.site_xvelp[self.ball_site_id, :].copy()
|
||||
elif self._steps == self.release_step and self.add_noise:
|
||||
self.sim.data.qvel[7::] += self.noise_std * np.random.randn(3)
|
||||
self.sim.data.qvel[7::] += self.noise_std * np.random.randn(3)
|
||||
crash = False
|
||||
except mujoco_py.builder.MujocoException:
|
||||
crash = True
|
||||
@@ -147,29 +141,32 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
|
||||
ball_pos = np.zeros(3)
|
||||
ball_vel = np.zeros(3)
|
||||
|
||||
infos = dict(reward_dist=reward_dist,
|
||||
reward=reward,
|
||||
velocity=angular_vel,
|
||||
# traj=self._q_pos,
|
||||
action=a,
|
||||
q_pos=self.sim.data.qpos[0:7].ravel().copy(),
|
||||
q_vel=self.sim.data.qvel[0:7].ravel().copy(),
|
||||
ball_pos=ball_pos,
|
||||
ball_vel=ball_vel,
|
||||
success=success,
|
||||
is_collided=is_collided, sim_crash=crash,
|
||||
table_contact_first=int(not self.reward_function.ball_ground_contact_first))
|
||||
infos = dict(
|
||||
reward_dist=reward_dist,
|
||||
reward=reward,
|
||||
velocity=angular_vel,
|
||||
# traj=self._q_pos,
|
||||
action=a,
|
||||
q_pos=self.sim.data.qpos[0:7].ravel().copy(),
|
||||
q_vel=self.sim.data.qvel[0:7].ravel().copy(),
|
||||
ball_pos=ball_pos,
|
||||
ball_vel=ball_vel,
|
||||
success=success,
|
||||
is_collided=is_collided, sim_crash=crash,
|
||||
table_contact_first=int(not self.reward_function.ball_ground_contact_first)
|
||||
)
|
||||
infos.update(reward_infos)
|
||||
return ob, reward, done, infos
|
||||
|
||||
def check_traj_in_joint_limits(self):
|
||||
def _check_traj_in_joint_limits(self):
|
||||
return any(self.current_pos > self.j_max) or any(self.current_pos < self.j_min)
|
||||
|
||||
def _get_obs(self):
|
||||
theta = self.sim.data.qpos.flat[:7]
|
||||
theta_dot = self.sim.data.qvel.flat[:7]
|
||||
ball_pos = self.sim.data.body_xpos[self.sim.model._body_name2id["ball"]].copy()
|
||||
cup_goal_diff_final = ball_pos - self.sim.data.site_xpos[self.sim.model._site_name2id["cup_goal_final_table"]].copy()
|
||||
cup_goal_diff_final = ball_pos - self.sim.data.site_xpos[
|
||||
self.sim.model._site_name2id["cup_goal_final_table"]].copy()
|
||||
cup_goal_diff_top = ball_pos - self.sim.data.site_xpos[self.sim.model._site_name2id["cup_goal_table"]].copy()
|
||||
return np.concatenate([
|
||||
np.cos(theta),
|
||||
@@ -179,17 +176,21 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
|
||||
cup_goal_diff_top,
|
||||
self.sim.model.body_pos[self.cup_table_id][:2].copy(),
|
||||
[self._steps],
|
||||
])
|
||||
])
|
||||
|
||||
def compute_reward(self):
|
||||
|
||||
@property
|
||||
def dt(self):
|
||||
return super(ALRBeerBongEnv, self).dt*self.repeat_action
|
||||
return super(ALRBeerBongEnv, self).dt * self.repeat_action
|
||||
|
||||
|
||||
class ALRBeerBongEnvFixedReleaseStep(ALRBeerBongEnv):
|
||||
def __init__(self, frame_skip=1, apply_gravity_comp=True, noisy=False, rndm_goal=False, cup_goal_pos=None):
|
||||
super().__init__(frame_skip, apply_gravity_comp, noisy, rndm_goal, cup_goal_pos)
|
||||
self.release_step = 62 # empirically evaluated for frame_skip=2!
|
||||
|
||||
|
||||
class ALRBeerBongEnvStepBasedEpisodicReward(ALRBeerBongEnv):
|
||||
def __init__(self, frame_skip=1, apply_gravity_comp=True, noisy=False, rndm_goal=False, cup_goal_pos=None):
|
||||
super().__init__(frame_skip, apply_gravity_comp, noisy, rndm_goal, cup_goal_pos)
|
||||
@@ -202,54 +203,21 @@ class ALRBeerBongEnvStepBasedEpisodicReward(ALRBeerBongEnv):
|
||||
reward = 0
|
||||
done = False
|
||||
while not done:
|
||||
sub_ob, sub_reward, done, sub_infos = super(ALRBeerBongEnvStepBasedEpisodicReward, self).step(np.zeros(a.shape))
|
||||
sub_ob, sub_reward, done, sub_infos = super(ALRBeerBongEnvStepBasedEpisodicReward, self).step(
|
||||
np.zeros(a.shape))
|
||||
reward += sub_reward
|
||||
infos = sub_infos
|
||||
ob = sub_ob
|
||||
ob[-1] = self.release_step + 1 # Since we simulate until the end of the episode, PPO does not see the
|
||||
# internal steps and thus, the observation also needs to be set correctly
|
||||
ob[-1] = self.release_step + 1 # Since we simulate until the end of the episode, PPO does not see the
|
||||
# internal steps and thus, the observation also needs to be set correctly
|
||||
return ob, reward, done, infos
|
||||
|
||||
|
||||
# class ALRBeerBongEnvStepBasedEpisodicReward(ALRBeerBongEnv):
|
||||
# def __init__(self, frame_skip=1, apply_gravity_comp=True, noisy=False, rndm_goal=False, cup_goal_pos=None):
|
||||
# super().__init__(frame_skip, apply_gravity_comp, noisy, rndm_goal, cup_goal_pos)
|
||||
# self.release_step = 62 # empirically evaluated for frame_skip=2!
|
||||
#
|
||||
# def step(self, a):
|
||||
# if self._steps < self.release_step:
|
||||
# return super(ALRBeerBongEnvStepBasedEpisodicReward, self).step(a)
|
||||
# else:
|
||||
# sub_ob, sub_reward, done, sub_infos = super(ALRBeerBongEnvStepBasedEpisodicReward, self).step(np.zeros(a.shape))
|
||||
# reward = sub_reward
|
||||
# infos = sub_infos
|
||||
# ob = sub_ob
|
||||
# ob[-1] = self.release_step + 1 # Since we simulate until the end of the episode, PPO does not see the
|
||||
# # internal steps and thus, the observation also needs to be set correctly
|
||||
# return ob, reward, done, infos
|
||||
|
||||
|
||||
class ALRBeerBongEnvStepBased(ALRBeerBongEnv):
|
||||
def __init__(self, frame_skip=1, apply_gravity_comp=True, noisy=False, rndm_goal=False, cup_goal_pos=None):
|
||||
super().__init__(frame_skip, apply_gravity_comp, noisy, rndm_goal, cup_goal_pos)
|
||||
self.release_step = 62 # empirically evaluated for frame_skip=2!
|
||||
|
||||
# def _set_action_space(self):
|
||||
# bounds = super(ALRBeerBongEnvStepBased, self)._set_action_space()
|
||||
# min_bound = np.concatenate(([-1], bounds.low), dtype=bounds.dtype)
|
||||
# max_bound = np.concatenate(([1], bounds.high), dtype=bounds.dtype)
|
||||
# self.action_space = spaces.Box(low=min_bound, high=max_bound, dtype=bounds.dtype)
|
||||
# return self.action_space
|
||||
|
||||
# def step(self, a):
|
||||
# self.release_step = self._steps if a[0]>=0 and self.release_step >= self._steps else self.release_step
|
||||
# return super(ALRBeerBongEnvStepBased, self).step(a[1:])
|
||||
#
|
||||
# def reset(self):
|
||||
# ob = super(ALRBeerBongEnvStepBased, self).reset()
|
||||
# self.release_step = self.ep_length + 1
|
||||
# return ob
|
||||
|
||||
def step(self, a):
|
||||
if self._steps < self.release_step:
|
||||
return super(ALRBeerBongEnvStepBased, self).step(a)
|
||||
@@ -267,9 +235,9 @@ class ALRBeerBongEnvStepBased(ALRBeerBongEnv):
|
||||
cup_goal_dist_top = np.linalg.norm(ball_pos - self.sim.data.site_xpos[
|
||||
self.sim.model._site_name2id["cup_goal_table"]].copy())
|
||||
if sub_infos['success']:
|
||||
dist_rew = -cup_goal_dist_final**2
|
||||
dist_rew = -cup_goal_dist_final ** 2
|
||||
else:
|
||||
dist_rew = -0.5*cup_goal_dist_final**2 - cup_goal_dist_top**2
|
||||
dist_rew = -0.5 * cup_goal_dist_final ** 2 - cup_goal_dist_top ** 2
|
||||
reward = reward - sub_infos['action_cost'] + dist_rew
|
||||
infos = sub_infos
|
||||
ob = sub_ob
|
||||
@@ -278,13 +246,13 @@ class ALRBeerBongEnvStepBased(ALRBeerBongEnv):
|
||||
return ob, reward, done, infos
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# env = ALRBeerBongEnv(rndm_goal=True)
|
||||
# env = ALRBeerBongEnvStepBased(frame_skip=2, rndm_goal=True)
|
||||
# env = ALRBeerBongEnvStepBasedEpisodicReward(frame_skip=2, rndm_goal=True)
|
||||
env = ALRBeerBongEnvFixedReleaseStep(frame_skip=2, rndm_goal=True)
|
||||
import time
|
||||
|
||||
env.reset()
|
||||
env.render("human")
|
||||
for i in range(1500):
|
||||
|
||||
@@ -29,52 +29,38 @@ class BeerPongReward:
|
||||
# "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)
|
||||
self.reset()
|
||||
self.is_initialized = False
|
||||
|
||||
def reset(self, noisy):
|
||||
self.ball_traj = []
|
||||
def reset(self):
|
||||
self.dists = []
|
||||
self.dists_final = []
|
||||
self.costs = []
|
||||
self.action_costs = []
|
||||
self.angle_rewards = []
|
||||
self.cup_angles = []
|
||||
self.cup_z_axes = []
|
||||
self.ball_ground_contact_first = False
|
||||
self.ball_table_contact = False
|
||||
self.ball_wall_contact = False
|
||||
self.ball_cup_contact = False
|
||||
self.ball_in_cup = False
|
||||
self.dist_ground_cup = -1 # distance floor to cup if first floor contact
|
||||
self.noisy_bp = noisy
|
||||
self._t_min_final_dist = -1
|
||||
self.dist_ground_cup = -1 # distance floor to cup if first floor contact
|
||||
|
||||
def initialize(self, env):
|
||||
|
||||
if not self.is_initialized:
|
||||
self.is_initialized = True
|
||||
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.ball_collision_id = env.sim.model._geom_name2id["ball_geom"]
|
||||
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.cup_collision_ids = [env.sim.model._geom_name2id[name] for name in self.cup_collision_objects]
|
||||
self.robot_collision_ids = [env.sim.model._geom_name2id[name] for name in self.robot_collision_objects]
|
||||
|
||||
def compute_reward(self, env, action):
|
||||
@@ -87,7 +73,7 @@ class BeerPongReward:
|
||||
self.check_contacts(env.sim)
|
||||
self.dists.append(np.linalg.norm(goal_pos - ball_pos))
|
||||
self.dists_final.append(np.linalg.norm(goal_final_pos - ball_pos))
|
||||
self.dist_ground_cup = np.linalg.norm(ball_pos-goal_pos) \
|
||||
self.dist_ground_cup = np.linalg.norm(ball_pos - goal_pos) \
|
||||
if self.ball_ground_contact_first and self.dist_ground_cup == -1 else self.dist_ground_cup
|
||||
action_cost = np.sum(np.square(action))
|
||||
self.action_costs.append(np.copy(action_cost))
|
||||
@@ -99,7 +85,7 @@ class BeerPongReward:
|
||||
min_dist = np.min(self.dists)
|
||||
final_dist = self.dists_final[-1]
|
||||
if self.ball_ground_contact_first:
|
||||
min_dist_coeff, final_dist_coeff, ground_contact_dist_coeff, rew_offset = 1, 0.5, 2, -4 # relative rew offset when first bounding on ground
|
||||
min_dist_coeff, final_dist_coeff, ground_contact_dist_coeff, rew_offset = 1, 0.5, 2, -4 # relative rew offset when first bounding on ground
|
||||
# min_dist_coeff, final_dist_coeff, ground_contact_dist_coeff, rew_offset = 1, 0.5, 0, -6 # absolute rew offset when first bouncing on ground
|
||||
else:
|
||||
if not self.ball_in_cup:
|
||||
@@ -108,18 +94,18 @@ class BeerPongReward:
|
||||
else:
|
||||
min_dist_coeff, final_dist_coeff, ground_contact_dist_coeff, rew_offset = 1, 0.5, 0, -2
|
||||
else:
|
||||
min_dist_coeff, final_dist_coeff, ground_contact_dist_coeff, rew_offset = 0, 1, 0 ,0
|
||||
min_dist_coeff, final_dist_coeff, ground_contact_dist_coeff, rew_offset = 0, 1, 0, 0
|
||||
# dist_ground_cup = 1 * self.dist_ground_cup
|
||||
action_cost = 1e-4 * np.mean(action_cost)
|
||||
reward = rew_offset - min_dist_coeff * min_dist ** 2 - final_dist_coeff * final_dist ** 2 - \
|
||||
action_cost - ground_contact_dist_coeff*self.dist_ground_cup ** 2
|
||||
action_cost - ground_contact_dist_coeff * self.dist_ground_cup ** 2
|
||||
# 1e-7*np.mean(action_cost)
|
||||
# release step punishment
|
||||
min_time_bound = 0.1
|
||||
max_time_bound = 1.0
|
||||
release_time = env.release_step*env.dt
|
||||
release_time_rew = int(release_time<min_time_bound)*(-30-10*(release_time-min_time_bound)**2) \
|
||||
+int(release_time>max_time_bound)*(-30-10*(release_time-max_time_bound)**2)
|
||||
release_time = env.release_step * env.dt
|
||||
release_time_rew = int(release_time < min_time_bound) * (-30 - 10 * (release_time - min_time_bound) ** 2) \
|
||||
+ int(release_time > max_time_bound) * (-30 - 10 * (release_time - max_time_bound) ** 2)
|
||||
reward += release_time_rew
|
||||
success = self.ball_in_cup
|
||||
# print('release time :', release_time)
|
||||
@@ -147,17 +133,17 @@ class BeerPongReward:
|
||||
self.table_collision_id)
|
||||
if not self.ball_cup_contact:
|
||||
self.ball_cup_contact = self._check_collision_with_set_of_objects(sim, self.ball_collision_id,
|
||||
self.cup_collision_ids)
|
||||
self.cup_collision_ids)
|
||||
if not self.ball_wall_contact:
|
||||
self.ball_wall_contact = self._check_collision_single_objects(sim, self.ball_collision_id,
|
||||
self.wall_collision_id)
|
||||
self.wall_collision_id)
|
||||
if not self.ball_in_cup:
|
||||
self.ball_in_cup = self._check_collision_single_objects(sim, self.ball_collision_id,
|
||||
self.cup_table_collision_id)
|
||||
if not self.ball_ground_contact_first:
|
||||
if not self.ball_table_contact and not self.ball_cup_contact and not self.ball_wall_contact and not self.ball_in_cup:
|
||||
self.ball_ground_contact_first = self._check_collision_single_objects(sim, self.ball_collision_id,
|
||||
self.ground_collision_id)
|
||||
self.ground_collision_id)
|
||||
|
||||
def _check_collision_single_objects(self, sim, id_1, id_2):
|
||||
for coni in range(0, sim.data.ncon):
|
||||
@@ -190,4 +176,4 @@ class BeerPongReward:
|
||||
|
||||
if collision or collision_trans:
|
||||
return True
|
||||
return False
|
||||
return False
|
||||
|
||||
@@ -1,10 +1,15 @@
|
||||
from alr_envs.mp.episodic_wrapper import EpisodicWrapper
|
||||
from typing import Union, Tuple
|
||||
from typing import Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import gym
|
||||
|
||||
from alr_envs.mp.episodic_wrapper import EpisodicWrapper
|
||||
|
||||
|
||||
class NewMPWrapper(EpisodicWrapper):
|
||||
|
||||
# def __init__(self, replanning_model):
|
||||
# self.replanning_model = replanning_model
|
||||
|
||||
@property
|
||||
def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
return self.env.sim.data.qpos[0:7].copy()
|
||||
@@ -22,26 +27,16 @@ class NewMPWrapper(EpisodicWrapper):
|
||||
[False] * 3, # cup_goal_diff_top
|
||||
[True] * 2, # xy position of cup
|
||||
[False] # env steps
|
||||
])
|
||||
])
|
||||
|
||||
# def set_mp_action_space(self):
|
||||
# min_action_bounds, max_action_bounds = self.mp.get_param_bounds()
|
||||
# if self.mp.learn_tau:
|
||||
# min_action_bounds[0] = 20*self.env.dt
|
||||
# max_action_bounds[0] = 260*self.env.dt
|
||||
# mp_action_space = gym.spaces.Box(low=min_action_bounds.numpy(), high=max_action_bounds.numpy(),
|
||||
# dtype=np.float32)
|
||||
# return mp_action_space
|
||||
|
||||
# def _step_callback(self, t: int, env_spec_params: Union[np.ndarray, None], step_action: np.ndarray) -> Union[np.ndarray]:
|
||||
# if self.mp.learn_tau:
|
||||
# return np.concatenate((step_action, np.atleast_1d(env_spec_params)))
|
||||
# else:
|
||||
# return step_action
|
||||
def do_replanning(self, pos, vel, s, a, t, last_replan_step):
|
||||
return False
|
||||
# const = np.arange(0, 1000, 10)
|
||||
# return bool(self.replanning_model(s))
|
||||
|
||||
def _episode_callback(self, action: np.ndarray) -> Tuple[np.ndarray, Union[np.ndarray, None]]:
|
||||
if self.mp.learn_tau:
|
||||
self.env.env.release_step = action[0]/self.env.dt # Tau value
|
||||
self.env.env.release_step = action[0] / self.env.dt # Tau value
|
||||
return action, None
|
||||
else:
|
||||
return action, None
|
||||
@@ -52,12 +47,3 @@ class NewMPWrapper(EpisodicWrapper):
|
||||
xyz[-1] = 0.840
|
||||
self.env.env.model.body_pos[self.env.env.cup_table_id] = xyz
|
||||
return self.get_observation_from_step(self.env.env._get_obs())
|
||||
# def set_action_space(self):
|
||||
# if self.mp.learn_tau:
|
||||
# min_action_bounds, max_action_bounds = self.mp.get_param_bounds()
|
||||
# min_action_bounds = np.concatenate((min_action_bounds.numpy(), [self.env.action_space.low[-1]]))
|
||||
# max_action_bounds = np.concatenate((max_action_bounds.numpy(), [self.env.action_space.high[-1]]))
|
||||
# self.action_space = gym.spaces.Box(low=min_action_bounds, high=max_action_bounds, dtype=np.float32)
|
||||
# return self.action_space
|
||||
# else:
|
||||
# return super(NewMPWrapper, self).set_action_space()
|
||||
|
||||
Reference in New Issue
Block a user