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

Merge branch 'master' into tt_cluster_debug

Browse Source 
# Conflicts:
#	fancy_gym/black_box/black_box_wrapper.py
#	fancy_gym/envs/__init__.py
#	fancy_gym/examples/examples_movement_primitives.py
This commit is contained in:
Hongyi Zhou
2022-11-25 17:16:26 +01:00
parent 932431a120 42863cee1c
commit c242c32a41
9 changed files with 326 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fancy_gym/black_box/black_box_wrapper.py
+25 -54
View File
@@ -2,7 +2,6 @@ from typing import Tuple, Optional, Callable
import gym
import numpy as np
import torch
from gym import spaces
from mp_pytorch.mp.mp_interfaces import MPInterface
@@ -23,8 +22,8 @@ class BlackBoxWrapper(gym.ObservationWrapper):
replanning_schedule: Optional[
Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int], bool]] = None,
reward_aggregation: Callable[[np.ndarray], float] = np.sum,
max_planning_times = None,
desired_conditioning: bool = False
max_planning_times: int = None,
condition_on_desired: bool = False
):
"""
gym.Wrapper for leveraging a black box approach with a trajectory generator.
@@ -59,10 +58,8 @@ class BlackBoxWrapper(gym.ObservationWrapper):
# reward computation
self.reward_aggregation = reward_aggregation
# self.traj_gen.basis_gn.show_basis(plot=True)
# spaces
self.return_context_observation = not (learn_sub_trajectories or self.do_replanning)
# self.return_context_observation = True
self.traj_gen_action_space = self._get_traj_gen_action_space()
self.action_space = self._get_action_space()
@@ -81,14 +78,12 @@ class BlackBoxWrapper(gym.ObservationWrapper):
self.verbose = verbose
# condition value
self.desired_conditioning = False
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.tau_first_prediction = None
self.plan_steps = 0
def observation(self, observation):
# return context space if we are
@@ -111,16 +106,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)
# self.traj_gen.set_duration(duration, self.dt)
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())
@@ -163,9 +153,6 @@ class BlackBoxWrapper(gym.ObservationWrapper):
def step(self, action: np.ndarray):
""" This function generates a trajectory based on a MP and then does the usual loop over reset and step"""
# time_is_valid = self.env.check_time_validity(action)
#
# if time_valid:
# TODO remove this part, right now only needed for beer pong
# mp_params, env_spec_params, proceed = self.env.episode_callback(action, self.traj_gen)
@@ -182,15 +169,13 @@ class BlackBoxWrapper(gym.ObservationWrapper):
infos = dict()
done = False
if self.verbose >= 2:
desired_pos_traj = []
desired_vel_traj = []
pos_traj = []
vel_traj = []
if traj_is_valid:
self.plan_counts += 1
self.plan_steps += 1
for t, (pos, vel) in enumerate(zip(position, velocity)):
current_pos = self.current_pos
current_vel = self.current_vel
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)
obs, c_reward, done, info = self.env.step(c_action)
@@ -205,42 +190,29 @@ class BlackBoxWrapper(gym.ObservationWrapper):
elems[t] = v
infos[k] = elems
if self.verbose >= 2:
desired_pos_traj.append(pos)
desired_vel_traj.append(vel)
pos_traj.append(self.current_pos)
vel_traj.append(self.current_vel)
if self.render_kwargs:
self.env.render(**self.render_kwargs)
if done or (self.replanning_schedule(self.current_pos, self.current_vel, obs, c_action,
t + 1 + self.current_traj_steps)
and self.max_planning_times is not None and self.plan_counts < self.max_planning_times):
if done or self.replanning_schedule(current_pos, 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:
# continue
if self.max_planning_times is not None and self.plan_steps >= self.max_planning_times:
continue
self.condition_pos = pos if self.desired_conditioning else self.current_pos
self.condition_vel = vel if self.desired_conditioning 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
infos.update({k: v[:t+1] for k, v in infos.items()})
self.current_traj_steps += t + 1
if self.verbose >= 2:
infos['desired_pos'] = position[:t+1]
infos['desired_vel'] = velocity[:t+1]
infos['current_pos'] = self.current_pos
infos['current_vel'] = self.current_vel
infos['step_actions'] = actions[:t + 1]
infos['step_observations'] = observations[:t + 1]
infos['step_rewards'] = rewards[:t + 1]
infos['desired_pos_traj'] = np.array(desired_pos_traj)
infos['desired_vel_traj'] = np.array(desired_vel_traj)
infos['pos_traj'] = np.array(pos_traj)
infos['vel_traj'] = np.array(vel_traj)
if self.verbose >= 2:
infos['positions'] = position
infos['velocities'] = velocity
infos['step_actions'] = actions[:t + 1]
infos['step_observations'] = observations[:t + 1]
infos['step_rewards'] = rewards[:t + 1]
infos['trajectory_length'] = t + 1
trajectory_return = self.reward_aggregation(rewards[:t + 1])
@@ -256,7 +228,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.tau_first_prediction = None
self.plan_steps = 0
self.traj_gen.reset()
return super(BlackBoxWrapper, self).reset()