make if condition better

This commit is contained in:
Hongyi Zhou 2022-11-28 10:49:05 +01:00
parent bb94c9c707
commit e7d2454e30

View File

@ -22,7 +22,7 @@ class BlackBoxWrapper(gym.ObservationWrapper):
replanning_schedule: Optional[ replanning_schedule: Optional[
Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int], bool]] = None, Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int], bool]] = None,
reward_aggregation: Callable[[np.ndarray], float] = np.sum, reward_aggregation: Callable[[np.ndarray], float] = np.sum,
max_planning_times: int = None, max_planning_times: int = np.inf,
condition_on_desired: bool = False condition_on_desired: bool = False
): ):
""" """
@ -178,11 +178,9 @@ class BlackBoxWrapper(gym.ObservationWrapper):
if self.render_kwargs: if self.render_kwargs:
self.env.render(**self.render_kwargs) self.env.render(**self.render_kwargs)
if done or self.replanning_schedule(self.current_pos, self.current_vel, obs, c_action, if done or (self.replanning_schedule(self.current_pos, self.current_vel, obs, c_action,
t + 1 + self.current_traj_steps): t + 1 + self.current_traj_steps)
and self.plan_steps < self.max_planning_times):
if not done and self.max_planning_times is not None and self.plan_steps >= self.max_planning_times:
continue
self.condition_pos = pos if self.condition_on_desired else None self.condition_pos = pos if self.condition_on_desired else None
self.condition_vel = vel if self.condition_on_desired else None self.condition_vel = vel if self.condition_on_desired else None