update according to reviews opinion & fix bugs in box pushing IK
This commit is contained in:
@@ -24,7 +24,7 @@ class BlackBoxWrapper(gym.ObservationWrapper):
|
||||
Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int], bool]] = None,
|
||||
reward_aggregation: Callable[[np.ndarray], float] = np.sum,
|
||||
max_planning_times: int = None,
|
||||
desired_traj_bc: bool = False
|
||||
condition_on_desired: bool = False
|
||||
):
|
||||
"""
|
||||
gym.Wrapper for leveraging a black box approach with a trajectory generator.
|
||||
@@ -71,12 +71,12 @@ class BlackBoxWrapper(gym.ObservationWrapper):
|
||||
self.verbose = verbose
|
||||
|
||||
# condition value
|
||||
self.desired_traj_bc = desired_traj_bc
|
||||
self.condition_on_desired = condition_on_desired
|
||||
self.condition_pos = None
|
||||
self.condition_vel = None
|
||||
|
||||
self.max_planning_times = max_planning_times
|
||||
self.plan_counts = 0
|
||||
self.plan_steps = 0
|
||||
|
||||
def observation(self, observation):
|
||||
# return context space if we are
|
||||
@@ -98,15 +98,11 @@ class BlackBoxWrapper(gym.ObservationWrapper):
|
||||
bc_time = np.array(0 if not self.do_replanning else self.current_traj_steps * self.dt)
|
||||
# TODO we could think about initializing with the previous desired value in order to have a smooth transition
|
||||
# at least from the planning point of view.
|
||||
# self.traj_gen.set_boundary_conditions(bc_time, self.current_pos, self.current_vel)
|
||||
if self.current_traj_steps == 0:
|
||||
self.condition_pos = self.current_pos
|
||||
self.condition_vel = self.current_vel
|
||||
|
||||
bc_time = torch.as_tensor(bc_time, dtype=torch.float32)
|
||||
self.condition_pos = torch.as_tensor(self.condition_pos, dtype=torch.float32)
|
||||
self.condition_vel = torch.as_tensor(self.condition_vel, dtype=torch.float32)
|
||||
self.traj_gen.set_boundary_conditions(bc_time, self.condition_pos, self.condition_vel)
|
||||
condition_pos = self.condition_pos if self.condition_pos is not None else self.current_pos
|
||||
condition_vel = self.condition_vel if self.condition_vel is not None else self.current_vel
|
||||
|
||||
self.traj_gen.set_boundary_conditions(bc_time, condition_pos, condition_vel)
|
||||
self.traj_gen.set_duration(duration, self.dt)
|
||||
# traj_dict = self.traj_gen.get_trajs(get_pos=True, get_vel=True)
|
||||
position = get_numpy(self.traj_gen.get_traj_pos())
|
||||
@@ -164,7 +160,7 @@ class BlackBoxWrapper(gym.ObservationWrapper):
|
||||
infos = dict()
|
||||
done = False
|
||||
|
||||
self.plan_counts += 1
|
||||
self.plan_steps += 1
|
||||
for t, (pos, vel) in enumerate(zip(position, velocity)):
|
||||
step_action = self.tracking_controller.get_action(pos, vel, self.current_pos, self.current_vel)
|
||||
c_action = np.clip(step_action, self.env.action_space.low, self.env.action_space.high)
|
||||
@@ -186,11 +182,11 @@ class BlackBoxWrapper(gym.ObservationWrapper):
|
||||
if done or self.replanning_schedule(self.current_pos, self.current_vel, obs, c_action,
|
||||
t + 1 + self.current_traj_steps):
|
||||
|
||||
if self.max_planning_times is not None and self.plan_counts >= self.max_planning_times:
|
||||
if self.max_planning_times is not None and self.plan_steps >= self.max_planning_times:
|
||||
continue
|
||||
|
||||
self.condition_pos = pos if self.desired_traj_bc else self.current_pos
|
||||
self.condition_vel = vel if self.desired_traj_bc else self.current_vel
|
||||
self.condition_pos = pos if self.condition_on_desired else None
|
||||
self.condition_vel = vel if self.condition_on_desired else None
|
||||
|
||||
break
|
||||
|
||||
@@ -215,6 +211,6 @@ class BlackBoxWrapper(gym.ObservationWrapper):
|
||||
|
||||
def reset(self, *, seed: Optional[int] = None, return_info: bool = False, options: Optional[dict] = None):
|
||||
self.current_traj_steps = 0
|
||||
self.plan_counts = 0
|
||||
self.plan_steps = 0
|
||||
self.traj_gen.reset()
|
||||
return super(BlackBoxWrapper, self).reset()
|
||||
|
||||
Reference in New Issue
Block a user