dodox/fancy_gym
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

updates

Browse Source
This commit is contained in:
Maximilian Huettenrauch
2021-04-23 11:37:42 +02:00
parent e482fc09f0
commit db001c411f
11 changed files with 207 additions and 186 deletions
Download Patch File Download Diff File Expand all files Collapse all files
alr_envs/mujoco/__init__.py
+1
View File
@@ -1,2 +1,3 @@
from alr_envs.mujoco.reacher.alr_reacher import ALRReacherEnv
from alr_envs.mujoco.balancing import BalancingEnv
from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup import ALRBallInACupEnv
alr_envs/mujoco/alr_mujoco_env.py
+17
View File
@@ -67,6 +67,9 @@ class AlrMujocoEnv(gym.Env):
self.init_qpos = self.sim.data.qpos.ravel().copy()
self.init_qvel = self.sim.data.qvel.ravel().copy()
self._start_pos = None
self._start_vel = None
self._set_action_space()
# action = self.action_space.sample()
@@ -93,6 +96,20 @@ class AlrMujocoEnv(gym.Env):
"""
return self.sim.data.qvel
@property
def start_pos(self):
"""
Start position of the agent, for example joint angles of a Panda robot. Necessary for MP wrapped envs.
"""
return self._start_pos
@property
def start_vel(self):
"""
Start velocity of the agent. Necessary for MP wrapped envs.
"""
return self._start_vel
def extend_des_pos(self, des_pos):
"""
In a simplified environment, the actions may only control a subset of all the joints.
alr_envs/mujoco/ball_in_a_cup/__init__.py
+1 -1
View File
@@ -1 +1 @@
from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_simple import ALRBallInACupEnv
alr_envs/mujoco/ball_in_a_cup/ball_in_a_cup.py
+25 -5
View File
@@ -5,15 +5,12 @@ from alr_envs.mujoco import alr_mujoco_env
class ALRBallInACupEnv(alr_mujoco_env.AlrMujocoEnv, utils.EzPickle):
def __init__(self, n_substeps=4, apply_gravity_comp=True, reward_function=None):
def __init__(self, n_substeps=4, apply_gravity_comp=True, reward_type: str = None, context: np.ndarray = None):
self._steps = 0
self.xml_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets",
"biac_base" + ".xml")
self.start_pos = np.array([0.0, 0.58760536, 0.0, 1.36004913, 0.0, -0.32072943, -1.57])
self.start_vel = np.zeros(7)
self._q_pos = []
self._q_vel = []
# self.weight_matrix_scale = 50
@@ -31,13 +28,21 @@ class ALRBallInACupEnv(alr_mujoco_env.AlrMujocoEnv, utils.EzPickle):
self.xml_path,
apply_gravity_comp=apply_gravity_comp,
n_substeps=n_substeps)
self._start_pos = np.array([0.0, 0.58760536, 0.0, 1.36004913, 0.0, -0.32072943, -1.57])
self._start_vel = np.zeros(7)
self.sim_time = 8 # seconds
self.sim_steps = int(self.sim_time / self.dt)
if reward_function is None:
if reward_type == "simple":
from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_reward_simple import BallInACupReward
reward_function = BallInACupReward
elif reward_type == "contextual_goal":
from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_reward import BallInACupReward
reward_function = BallInACupReward
else:
raise ValueError("Unknown reward type")
self.reward_function = reward_function(self.sim_steps)
self.configure(context)
@property
def current_pos(self):
@@ -97,6 +102,7 @@ class ALRBallInACupEnv(alr_mujoco_env.AlrMujocoEnv, utils.EzPickle):
def check_traj_in_joint_limits(self):
return any(self.current_pos > self.j_max) or any(self.current_pos < self.j_min)
# TODO: extend observation space
def _get_obs(self):
theta = self.sim.data.qpos.flat[:7]
return np.concatenate([
@@ -106,6 +112,20 @@ class ALRBallInACupEnv(alr_mujoco_env.AlrMujocoEnv, utils.EzPickle):
[self._steps],
])
# These functions are for the task with 3 joint actuations
def extend_des_pos(self, des_pos):
des_pos_full = self.start_pos.copy()
des_pos_full[1] = des_pos[0]
des_pos_full[3] = des_pos[1]
des_pos_full[5] = des_pos[2]
return des_pos_full
def extend_des_vel(self, des_vel):
des_vel_full = self.start_vel.copy()
des_vel_full[1] = des_vel[0]
des_vel_full[3] = des_vel[1]
des_vel_full[5] = des_vel[2]
return des_vel_full
if __name__ == "__main__":
alr_envs/mujoco/ball_in_a_cup/ball_in_a_cup_reward.py
+23 -13
View File
@@ -40,9 +40,12 @@ class BallInACupReward(alr_reward_fct.AlrReward):
self.ball_in_cup = False
self.ball_above_threshold = False
self.dist_ctxt = 3
self.action_costs = []
self.cup_angles = []
def compute_reward(self, action, sim, step):
action_cost = np.sum(np.square(action))
self.action_costs.append(action_cost)
stop_sim = False
success = False
@@ -54,7 +57,7 @@ class BallInACupReward(alr_reward_fct.AlrReward):
self.collision_ids = [sim.model._geom_name2id[name] for name in self.collision_objects]
if self.check_collision(sim):
reward = - 1e-4 * action_cost - 1000
reward = - 1e-3 * action_cost - 1000
stop_sim = True
return reward, success, stop_sim
@@ -65,8 +68,11 @@ class BallInACupReward(alr_reward_fct.AlrReward):
self.dists.append(np.linalg.norm(goal_pos - ball_pos))
self.dists_final.append(np.linalg.norm(goal_final_pos - ball_pos))
self.dists_ctxt.append(np.linalg.norm(ball_pos - self.context))
self.ball_traj[step, :] = ball_pos
self.cup_traj[step, :] = goal_pos
self.ball_traj[step, :] = np.copy(ball_pos)
self.cup_traj[step, :] = np.copy(goal_pos) # ?
cup_quat = np.copy(sim.data.body_xquat[sim.model._body_name2id["cup"]])
self.cup_angles.append(np.arctan2(2 * (cup_quat[0] * cup_quat[1] + cup_quat[2] * cup_quat[3]),
1 - 2 * (cup_quat[1] ** 2 + cup_quat[2] ** 2)))
# Determine the first time when ball is in cup
if not self.ball_in_cup:
@@ -77,25 +83,29 @@ class BallInACupReward(alr_reward_fct.AlrReward):
self.dist_ctxt = dist_to_ctxt
if step == self.sim_time - 1:
t_min_dist = np.argmin(self.dists)
angle_min_dist = self.cup_angles[t_min_dist]
cost_angle = (angle_min_dist - np.pi / 2) ** 2
min_dist = np.min(self.dists)
dist_final = self.dists_final[-1]
# dist_ctxt = self.dists_ctxt[-1]
# max distance between ball and cup and cup height at that time
ball_to_cup_diff = self.ball_traj[:, 2] - self.cup_traj[:, 2]
t_max_diff = np.argmax(ball_to_cup_diff)
t_max_ball_height = np.argmax(self.ball_traj[:, 2])
max_ball_height = np.max(self.ball_traj[:, 2])
# # max distance between ball and cup and cup height at that time
# ball_to_cup_diff = self.ball_traj[:, 2] - self.cup_traj[:, 2]
# t_max_diff = np.argmax(ball_to_cup_diff)
# t_max_ball_height = np.argmax(self.ball_traj[:, 2])
# max_ball_height = np.max(self.ball_traj[:, 2])
# cost = self._get_stage_wise_cost(ball_in_cup, min_dist, dist_final, dist_ctxt)
cost = 0.3 * min_dist + 0.3 * dist_final + 0.3 * np.minimum(self.dist_ctxt, 3)
reward = np.exp(-1 * cost) - 1e-4 * action_cost
if max_ball_height < self.context[2] or ball_to_cup_diff[t_max_ball_height] < 0:
reward -= 1
cost = 0.5 * min_dist + 0.5 * dist_final + 0.3 * np.minimum(self.dist_ctxt, 3) + 0.01 * cost_angle
reward = np.exp(-2 * cost) - 1e-3 * action_cost
# if max_ball_height < self.context[2] or ball_to_cup_diff[t_max_ball_height] < 0:
# reward -= 1
success = dist_final < 0.05 and self.dist_ctxt < 0.05
else:
reward = - 1e-4 * action_cost
reward = - 1e-3 * action_cost
success = False
return reward, success, stop_sim
alr_envs/mujoco/ball_in_a_cup/ball_in_a_cup_simple.py
-151
View File
@@ -1,151 +0,0 @@
from alr_envs.mujoco import alr_mujoco_env
from gym import utils
import os
import numpy as np
class ALRBallInACupEnv(alr_mujoco_env.AlrMujocoEnv, utils.EzPickle):
def __init__(self, n_substeps=4, apply_gravity_comp=True, reward_function=None):
self._steps = 0
self.xml_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets",
"biac_base" + ".xml")
self.start_pos = np.array([0.0, 0.58760536, 0.0, 1.36004913, 0.0, -0.32072943, -1.57])
self.start_vel = np.zeros(7)
self._q_pos = []
self._q_vel = []
# self.weight_matrix_scale = 50
self.max_ctrl = np.array([150., 125., 40., 60., 5., 5., 2.])
self.p_gains = 1 / self.max_ctrl * np.array([200, 300, 100, 100, 10, 10, 2.5])
self.d_gains = 1 / self.max_ctrl * np.array([7, 15, 5, 2.5, 0.3, 0.3, 0.05])
self.j_min = np.array([-2.6, -1.985, -2.8, -0.9, -4.55, -1.5707, -2.7])
self.j_max = np.array([2.6, 1.985, 2.8, 3.14159, 1.25, 1.5707, 2.7])
self.context = None
utils.EzPickle.__init__(self)
alr_mujoco_env.AlrMujocoEnv.__init__(self,
self.xml_path,
apply_gravity_comp=apply_gravity_comp,
n_substeps=n_substeps)
self.sim_time = 8 # seconds
self.sim_steps = int(self.sim_time / self.dt)
if reward_function is None:
from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_reward_simple import BallInACupReward
reward_function = BallInACupReward
self.reward_function = reward_function(self.sim_steps)
@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 configure(self, context):
self.context = context
self.reward_function.reset(context)
def reset_model(self):
init_pos_all = self.init_qpos.copy()
init_pos_robot = self.start_pos
init_vel = np.zeros_like(init_pos_all)
ball_id = self.sim.model._body_name2id["ball"]
goal_final_id = self.sim.model._site_name2id["cup_goal_final"]
self._steps = 0
self._q_pos = []
self._q_vel = []
start_pos = init_pos_all
start_pos[0:7] = init_pos_robot
self.set_state(start_pos, init_vel)
return self._get_obs()
def step(self, a):
reward_dist = 0.0
angular_vel = 0.0
reward_ctrl = - np.square(a).sum()
crash = self.do_simulation(a)
joint_cons_viol = self.check_traj_in_joint_limits()
self._q_pos.append(self.sim.data.qpos[0:7].ravel().copy())
self._q_vel.append(self.sim.data.qvel[0:7].ravel().copy())
ob = self._get_obs()
if not crash and not joint_cons_viol:
reward, success, collision = self.reward_function.compute_reward(a, self.sim, self._steps)
done = success or self._steps == self.sim_steps - 1 or collision
self._steps += 1
else:
reward = -1000
success = False
done = True
return ob, reward, done, dict(reward_dist=reward_dist,
reward_ctrl=reward_ctrl,
velocity=angular_vel,
traj=self._q_pos, is_success=success,
is_collided=crash or joint_cons_viol)
def check_traj_in_joint_limits(self):
return any(self.current_pos > self.j_max) or any(self.current_pos < self.j_min)
# TODO: extend observation space
def _get_obs(self):
theta = self.sim.data.qpos.flat[:7]
return np.concatenate([
np.cos(theta),
np.sin(theta),
# self.get_body_com("target"), # only return target to make problem harder
[self._steps],
])
def extend_des_pos(self, des_pos):
des_pos_full = self.start_pos.copy()
des_pos_full[1] = des_pos[0]
des_pos_full[3] = des_pos[1]
des_pos_full[5] = des_pos[2]
return des_pos_full
def extend_des_vel(self, des_vel):
des_vel_full = self.start_vel.copy()
des_vel_full[1] = des_vel[0]
des_vel_full[3] = des_vel[1]
des_vel_full[5] = des_vel[2]
return des_vel_full
if __name__ == "__main__":
from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_reward_simple import BallInACupReward
env = ALRBallInACupEnv(reward_function=BallInACupReward)
env.configure(None)
env.reset()
env.render()
for i in range(4000):
# objective.load_result("/tmp/cma")
# test with random actions
# ac = 0.1 * env.action_space.sample()
# ac = -np.array([i, i, i]) / 10000 + np.array([env.start_pos[1], env.start_pos[3], env.start_pos[5]])
# ac = np.array([0., -10, 0, -1, 0, 1, 0])
ac = np.array([0.,0., 0, 0, 0, 0, 0])
# ac[0] += np.pi/2
obs, rew, d, info = env.step(ac)
env.render()
print(rew)
if d:
break
env.close()
alr_envs/mujoco/ball_in_a_cup/utils.py
+4 -5
View File
@@ -1,7 +1,6 @@
from alr_envs.utils.wrapper.detpmp_wrapper import DetPMPWrapper
from alr_envs.utils.wrapper.dmp_wrapper import DmpWrapper
from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup import ALRBallInACupEnv
from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_simple import ALRBallInACupEnv as ALRBallInACupEnvSimple
def make_contextual_env(rank, seed=0):
@@ -16,7 +15,7 @@ def make_contextual_env(rank, seed=0):
"""
def _init():
env = ALRBallInACupEnv()
env = ALRBallInACupEnv(reward_type="contextual_goal")
env = DetPMPWrapper(env,
num_dof=7,
@@ -50,7 +49,7 @@ def make_env(rank, seed=0):
"""
def _init():
env = ALRBallInACupEnvSimple()
env = ALRBallInACupEnv(reward_type="simple")
env = DetPMPWrapper(env,
num_dof=7,
@@ -84,7 +83,7 @@ def make_simple_env(rank, seed=0):
"""
def _init():
env = ALRBallInACupEnvSimple()
env = ALRBallInACupEnv(reward_type="simple")
env = DetPMPWrapper(env,
num_dof=3,
@@ -119,7 +118,7 @@ def make_simple_dmp_env(rank, seed=0):
"""
def _init():
_env = ALRBallInACupEnvSimple()
_env = ALRBallInACupEnv(reward_type="simple")
_env = DmpWrapper(_env,
num_dof=3,