working bp version

2022-05-03 19:51:54 +02:00
parent f33996c27a
commit 2fbde9fbb1
21 changed files with 460 additions and 224 deletions
@@ -1 +1,2 @@
-from .mp_wrapper import MPWrapper
+from .mp_wrapper import MPWrapper
+from .new_mp_wrapper import NewMPWrapper
@@ -3,6 +3,7 @@ import os

 import numpy as np
 from gym import utils
+from gym import spaces
 from gym.envs.mujoco import MujocoEnv
 from alr_envs.alr.mujoco.beerpong.beerpong_reward_staged import BeerPongReward

@@ -17,7 +18,7 @@ CUP_POS_MAX = np.array([0.32, -1.2])

 class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
    def __init__(self, frame_skip=1, apply_gravity_comp=True, noisy=False,
-                 rndm_goal=False, cup_goal_pos=None):
+                 rndm_goal=False, learn_release_step=True, 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:
            cup_goal_pos = np.insert(cup_goal_pos, 2, 0.840)
@@ -51,10 +52,9 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
            self.noise_std = 0.01
        else:
            self.noise_std = 0
-
+        self.learn_release_step = learn_release_step
        reward_function = BeerPongReward
        self.reward_function = reward_function()
-        self.n_table_bounces_first = 0

        MujocoEnv.__init__(self, self.xml_path, frame_skip)
        utils.EzPickle.__init__(self)
@@ -63,6 +63,13 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
    def start_pos(self):
        return self._start_pos

+    def _set_action_space(self):
+        bounds = self.model.actuator_ctrlrange.copy().astype(np.float32)
+        bounds = np.concatenate((bounds, [[50, self.ep_length*0.333]]), axis=0)
+        low, high = bounds.T
+        self.action_space = spaces.Box(low=low, high=high, dtype=np.float32)
+        return self.action_space
+
    @property
    def start_vel(self):
        return self._start_vel
@@ -76,8 +83,6 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
        return self.sim.data.qvel[0:7].copy()

    def reset(self):
-        print(not self.reward_function.ball_ground_contact_first)
-        self.n_table_bounces_first += int(not self.reward_function.ball_ground_contact_first)
        self.reward_function.reset(self.add_noise)
        return super().reset()

@@ -104,14 +109,17 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
        return self._get_obs()

    def step(self, a):
+        self._release_step = a[-1] if self.learn_release_step else self._release_step
+        self._release_step = np.clip(self._release_step, self.action_space.low[-1], self.action_space.high[-1]) \
+            if self.learn_release_step else self._release_step
        reward_dist = 0.0
        angular_vel = 0.0
-        reward_ctrl = - np.square(a).sum()
-
+        applied_action = a[:a.shape[0]-int(self.learn_release_step)]
+        reward_ctrl = - np.square(applied_action).sum()
        if self.apply_gravity_comp:
-            a = a + self.sim.data.qfrc_bias[:len(a)].copy() / self.model.actuator_gear[:, 0]
+            applied_action += self.sim.data.qfrc_bias[:len(applied_action)].copy() / self.model.actuator_gear[:, 0]
        try:
-            self.do_simulation(a, self.frame_skip)
+            self.do_simulation(applied_action, self.frame_skip)
            if self._steps < self._release_step:
                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()
@@ -125,7 +133,7 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
        ob = self._get_obs()

        if not crash:
-            reward, reward_infos = self.reward_function.compute_reward(self, a)
+            reward, reward_infos = self.reward_function.compute_reward(self, applied_action)
            success = reward_infos['success']
            is_collided = reward_infos['is_collided']
            ball_pos = reward_infos['ball_pos']
@@ -162,6 +162,10 @@ class BeerPongReward:
                min_dist_coeff, final_dist_coeff, rew_offset = 0, 1, 0
            reward = rew_offset - min_dist_coeff * min_dist ** 2 - final_dist_coeff * final_dist ** 2 - \
                     1e-4 * np.mean(action_cost)
+            if env.learn_release_step and not self.ball_in_cup:
+                too_small = (env._release_step<50)*(env._release_step-50)**2
+                too_big = (env._release_step>200)*0.2*(env._release_step-200)**2
+                reward = reward - too_small -too_big
            # 1e-7*np.mean(action_cost)
            success = self.ball_in_cup
        else:
@@ -1,13 +1,15 @@
-from mp_wrapper import BaseMPWrapper
+from alr_envs.mp.episodic_wrapper import EpisodicWrapper
 from typing import Union, Tuple
 import numpy as np
+import gym


-class MPWrapper(BaseMPWrapper):
-
+class NewMPWrapper(EpisodicWrapper):
+    @property
    def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
        return self.env.sim.data.qpos[0:7].copy()

+    @property
    def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
        return self.env.sim.data.qvel[0:7].copy()

@@ -18,3 +20,25 @@ class MPWrapper(BaseMPWrapper):
            [True] * 2,  # xy position of cup
            [False]  # env steps
        ])
+
+    def _step_callback(self, t: int, env_spec_params: Union[np.ndarray, None], step_action: np.ndarray) -> Union[np.ndarray]:
+        if self.env.learn_release_step:
+            return np.concatenate((step_action, np.atleast_1d(env_spec_params)))
+        else:
+            return step_action
+
+    def _episode_callback(self, action: np.ndarray) -> Tuple[np.ndarray, Union[np.ndarray, None]]:
+        if self.env.learn_release_step:
+            return action[:-1], action[-1]  # mp_params, release step
+        else:
+            return action, None
+
+    def set_action_space(self):
+        if self.env.learn_release_step:
+            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.mp_action_space = gym.spaces.Box(low=min_action_bounds, high=max_action_bounds, dtype=np.float32)
+            return self.mp_action_space
+        else:
+            return super(NewMPWrapper, self).set_action_space()
@@ -1,9 +1,9 @@
-from mp_wrapper import BaseMPWrapper
+from alr_envs.mp.episodic_wrapper import EpisodicWrapper
 from typing import Union, Tuple
 import numpy as np


-class MPWrapper(BaseMPWrapper):
+class MPWrapper(EpisodicWrapper):

    @property
    def current_pos(self) -> Union[float, int, np.ndarray, Tuple]: