fix minor bug

This commit is contained in:
Hongyi Zhou 2022-12-01 13:16:37 +01:00
parent a9a1d05497
commit 55df1e0ef6

View File

@ -147,7 +147,7 @@ class BlackBoxWrapper(gym.ObservationWrapper):
position, velocity = self.get_trajectory(action) position, velocity = self.get_trajectory(action)
position, velocity = self.env.set_episode_arguments(action, position, velocity) position, velocity = self.env.set_episode_arguments(action, position, velocity)
traj_is_valid = self.env.preprocessing_and_validity_callback(action, position, velocity) traj_is_valid, position, velocity = self.env.preprocessing_and_validity_callback(action, position, velocity)
trajectory_length = len(position) trajectory_length = len(position)
rewards = np.zeros(shape=(trajectory_length,)) rewards = np.zeros(shape=(trajectory_length,))
@ -159,7 +159,7 @@ class BlackBoxWrapper(gym.ObservationWrapper):
infos = dict() infos = dict()
done = False done = False
if not traj_is_valid: if traj_is_valid is False:
obs, trajectory_return, done, infos = self.env.invalid_traj_callback(action, position, velocity, obs, trajectory_return, done, infos = self.env.invalid_traj_callback(action, position, velocity,
self.return_context_observation) self.return_context_observation)
return self.observation(obs), trajectory_return, done, infos return self.observation(obs), trajectory_return, done, infos