class VectorizedSampler(BaseSampler):
def __init__(self, algo, n_envs=None):
super(VectorizedSampler, self).__init__(algo)
self.n_envs = n_envs
def start_worker(self):
n_envs = self.n_envs
if n_envs is None:
n_envs = int(self.algo.batch_size / self.algo.max_path_length)
n_envs = max(1, min(n_envs, 100))
if getattr(self.algo.env, 'vectorized', False):
self.vec_env = self.algo.env.vec_env_executor(
n_envs=n_envs, max_path_length=self.algo.max_path_length)
else:
envs = [
pickle.loads(pickle.dumps(self.algo.env))
for _ in range(n_envs)
]
self.vec_env = VecEnvExecutor(
envs=envs, max_path_length=self.algo.max_path_length)
self.env_spec = self.algo.env.spec
def shutdown_worker(self):
self.vec_env.terminate()
def obtain_samples(self, itr):
logger.log("Obtaining samples for iteration %d..." % itr)
paths = []
n_samples = 0
obses = self.vec_env.reset()
dones = np.asarray([True] * self.vec_env.num_envs)
running_paths = [None] * self.vec_env.num_envs
pbar = ProgBarCounter(self.algo.batch_size)
policy_time = 0
env_time = 0
process_time = 0
policy = self.algo.policy
import time
while n_samples < self.algo.batch_size:
t = time.time()
policy.reset(dones)
actions, agent_infos = policy.get_actions(obses)
policy_time += time.time() - t
t = time.time()
next_obses, rewards, dones, env_infos = self.vec_env.step(actions)
env_time += time.time() - t
t = time.time()
agent_infos = tensor_utils.split_tensor_dict_list(agent_infos)
env_infos = tensor_utils.split_tensor_dict_list(env_infos)
if env_infos is None:
env_infos = [dict() for _ in range(self.vec_env.num_envs)]
if agent_infos is None:
agent_infos = [dict() for _ in range(self.vec_env.num_envs)]
for idx, observation, action, reward, env_info, agent_info, done in zip(
itertools.count(), obses, actions, rewards, env_infos,
agent_infos, dones):
if running_paths[idx] is None:
running_paths[idx] = dict(
observations=[],
actions=[],
rewards=[],
terminals=[],
env_infos=[],
agent_infos=[],
)
running_paths[idx]["observations"].append(observation)
running_paths[idx]["actions"].append(action)
running_paths[idx]["rewards"].append(reward)
running_paths[idx]["terminals"].append(done)
running_paths[idx]["env_infos"].append(env_info)
running_paths[idx]["agent_infos"].append(agent_info)
if done:
paths.append(
dict(
observations=self.env_spec.observation_space.
flatten_n(running_paths[idx]["observations"]),
actions=self.env_spec.action_space.flatten_n(
running_paths[idx]["actions"]),
rewards=tensor_utils.stack_tensor_list(
running_paths[idx]["rewards"]),
terminals=tensor_utils.stack_tensor_list(
running_paths[idx]["terminals"]),
env_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["env_infos"]),
agent_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["agent_infos"]),
))
n_samples += len(running_paths[idx]["rewards"])
running_paths[idx] = None
process_time += time.time() - t
pbar.inc(len(obses))
obses = next_obses
pbar.stop()
logger.record_tabular("PolicyExecTime", policy_time)
logger.record_tabular("EnvExecTime", env_time)
logger.record_tabular("ProcessExecTime", process_time)
return paths
class VectorizedSampler(BaseSampler):
def __init__(self, algo, n_envs=None):
super(VectorizedSampler, self).__init__(algo)
self.n_envs = n_envs
def start_worker(self):
n_envs = self.n_envs
if n_envs is None:
n_envs = int(self.algo.batch_size / self.algo.max_path_length)
n_envs = max(1, min(n_envs, 100))
if getattr(self.algo.env, 'vectorized', False):
self.vec_env = self.algo.env.vec_env_executor(
n_envs=n_envs, max_path_length=self.algo.max_path_length)
else:
envs = [
pickle.loads(pickle.dumps(self.algo.env))
for _ in range(n_envs)
]
self.vec_env = VecEnvExecutor(
envs=envs,
#env=pickle.loads(pickle.dumps(self.algo.env)),
#n = n_envs,
max_path_length=self.algo.max_path_length)
self.env_spec = self.algo.env.spec
def shutdown_worker(self):
self.vec_env.terminate()
def obtain_samples(self,
itr,
reset_args=None,
return_dict=False,
log_prefix=''):
# reset_args: arguments to pass to the environments to reset
# return_dict: whether or not to return a dictionary or list form of paths
logger.log("Obtaining samples for iteration %d..." % itr)
#paths = []
paths = {}
for i in range(self.vec_env.num_envs):
paths[i] = []
# if the reset args are not list/numpy, we set the same args for each env
if reset_args is not None and (type(reset_args) != list
and type(reset_args) != np.ndarray):
reset_args = [reset_args] * self.vec_env.num_envs
n_samples = 0
obses = self.vec_env.reset(reset_args)
dones = np.asarray([True] * self.vec_env.num_envs)
running_paths = [None] * self.vec_env.num_envs
pbar = ProgBarCounter(self.algo.batch_size)
policy_time = 0
env_time = 0
process_time = 0
policy = self.algo.policy
import time
while n_samples < self.algo.batch_size:
t = time.time()
policy.reset(dones)
actions, agent_infos = policy.get_actions(obses)
policy_time += time.time() - t
t = time.time()
next_obses, rewards, dones, env_infos = self.vec_env.step(
actions, reset_args)
env_time += time.time() - t
t = time.time()
agent_infos = tensor_utils.split_tensor_dict_list(agent_infos)
env_infos = tensor_utils.split_tensor_dict_list(env_infos)
if env_infos is None:
env_infos = [dict() for _ in range(self.vec_env.num_envs)]
if agent_infos is None:
agent_infos = [dict() for _ in range(self.vec_env.num_envs)]
for idx, observation, action, reward, env_info, agent_info, done in zip(
itertools.count(), obses, actions, rewards, env_infos,
agent_infos, dones):
if running_paths[idx] is None:
running_paths[idx] = dict(
observations=[],
actions=[],
rewards=[],
env_infos=[],
agent_infos=[],
)
running_paths[idx]["observations"].append(observation)
running_paths[idx]["actions"].append(action)
running_paths[idx]["rewards"].append(reward)
running_paths[idx]["env_infos"].append(env_info)
running_paths[idx]["agent_infos"].append(agent_info)
if done:
paths[idx].append(
dict(
observations=self.env_spec.observation_space.
flatten_n(running_paths[idx]["observations"]),
actions=self.env_spec.action_space.flatten_n(
running_paths[idx]["actions"]),
rewards=tensor_utils.stack_tensor_list(
running_paths[idx]["rewards"]),
env_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["env_infos"]),
agent_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["agent_infos"]),
))
n_samples += len(running_paths[idx]["rewards"])
running_paths[idx] = None
process_time += time.time() - t
pbar.inc(len(obses))
obses = next_obses
pbar.stop()
#logger.record_tabular(log_prefix+"PolicyExecTime", policy_time)
#logger.record_tabular(log_prefix+"EnvExecTime", env_time)
#logger.record_tabular(log_prefix+"ProcessExecTime", process_time)
if not return_dict:
flatten_list = lambda l: [
item for sublist in l for item in sublist
]
paths = flatten_list(paths.values())
#path_keys = flatten_list([[key]*len(paths[key]) for key in paths.keys()])
return paths
class VectorizedSampler(BaseSampler):
def start_worker(self):
estimated_envs = int(self.algo.batch_size / self.algo.max_path_length)
estimated_envs = max(1, min(estimated_envs, 100))
self.vec_env = VecEnvExecutor(
self.algo.env,
n=estimated_envs,
max_path_length=self.algo.max_path_length
)
self.env_spec = self.algo.env.spec
def shutdown_worker(self):
self.vec_env.terminate()
def obtain_samples(self, itr):
logger.log("Obtaining samples for iteration %d..." % itr)
paths = []
n_samples = 0
obses = self.vec_env.reset()
dones = np.asarray([True] * self.vec_env.num_envs)
running_paths = [None] * self.vec_env.num_envs
pbar = ProgBarCounter(self.algo.batch_size)
policy_time = 0
env_time = 0
process_time = 0
import time
while n_samples < self.algo.batch_size:
t = time.time()
self.algo.policy.reset(dones)
actions, agent_infos = self.algo.policy.get_actions(obses)
policy_time += time.time() - t
t = time.time()
next_obses, rewards, dones, env_infos = self.vec_env.step(actions)
env_time += time.time() - t
t = time.time()
agent_infos = tensor_utils.split_tensor_dict_list(agent_infos)
env_infos = tensor_utils.split_tensor_dict_list(env_infos)
if env_infos is None:
env_infos = [dict() for _ in xrange(self.vec_env.num_envs)]
if agent_infos is None:
agent_infos = [dict() for _ in xrange(self.vec_env.num_envs)]
for idx, observation, action, reward, env_info, agent_info, done in zip(itertools.count(), obses, actions,
rewards, env_infos, agent_infos,
dones):
if running_paths[idx] is None:
running_paths[idx] = dict(
observations=[],
actions=[],
rewards=[],
env_infos=[],
agent_infos=[],
)
running_paths[idx]["observations"].append(observation)
running_paths[idx]["actions"].append(action)
running_paths[idx]["rewards"].append(reward)
running_paths[idx]["env_infos"].append(env_info)
running_paths[idx]["agent_infos"].append(agent_info)
if done:
paths.append(dict(
observations=self.env_spec.observation_space.flatten_n(running_paths[idx]["observations"]),
actions=self.env_spec.action_space.flatten_n(running_paths[idx]["actions"]),
rewards=tensor_utils.stack_tensor_list(running_paths[idx]["rewards"]),
env_infos=tensor_utils.stack_tensor_dict_list(running_paths[idx]["env_infos"]),
agent_infos=tensor_utils.stack_tensor_dict_list(running_paths[idx]["agent_infos"]),
))
n_samples += len(running_paths[idx]["rewards"])
running_paths[idx] = None
process_time += time.time() - t
pbar.inc(len(obses))
obses = next_obses
pbar.stop()
logger.record_tabular("PolicyExecTime", policy_time)
logger.record_tabular("EnvExecTime", env_time)
logger.record_tabular("ProcessExecTime", process_time)
return paths
class VectorizedSampler(BaseSampler):
def __init__(self, algo, n_envs=None, batch_size=None):
super(VectorizedSampler, self).__init__(algo)
self.n_envs = n_envs
# if batch_size is not None:
# self.batch_size = batch_size
# else:
self.batch_size = self.algo.batch_size
print("vectorized sampler initiated")
def start_worker(self):
n_envs = self.n_envs
if n_envs is None:
n_envs = int(self.algo.batch_size / self.algo.max_path_length)
n_envs = max(1, min(n_envs, 100))
if getattr(self.algo.env, 'vectorized', False):
self.vec_env = self.algo.env.vec_env_executor(n_envs=n_envs, max_path_length=self.algo.max_path_length)
else:
envs = [pickle.loads(pickle.dumps(self.algo.env)) for _ in range(n_envs)]
self.vec_env = VecEnvExecutor(
envs=envs,
#env=pickle.loads(pickle.dumps(self.algo.env)),
#n = n_envs,
max_path_length=self.algo.max_path_length
)
self.env_spec = self.algo.env.spec
def shutdown_worker(self):
self.vec_env.terminate()
def obtain_samples(self, itr, reset_args=None, return_dict=False, log_prefix='', preupdate=False, save_img_obs=False, contexts = None):
# reset_args: arguments to pass to the environments to reset
# return_dict: whether or not to return a dictionary or list form of paths
logger.log("Obtaining samples for iteration %d..." % itr)
paths = {}
for i in range(self.vec_env.num_envs):
paths[i] = []
# if the reset args are not list/numpy, we set the same args for each env
if reset_args is not None and (type(reset_args) != list and type(reset_args) != np.ndarray):
assert False, "debug, should we be using this?"
print("WARNING, will vectorize reset_args")
reset_args = [reset_args]*self.vec_env.num_envs
n_samples = 0
path_nums = [0] * self.vec_env.num_envs # keeps track on which rollout we are for each environment instance
obses = self.vec_env.reset(reset_args)
if contexts:
obses = np.concatenate([obses, contexts], axis = 1)
dones = np.asarray([True] * self.vec_env.num_envs)
running_paths = [None] * self.vec_env.num_envs
pbar = ProgBarCounter(self.batch_size)
policy_time = 0
env_time = 0
process_time = 0
if contexts:
policy = self.algo.post_policy
else:
policy = self.algo.policy
while n_samples < self.batch_size:
t = time.time()
policy.reset(dones)
actions, agent_infos = policy.get_actions(obses)
# print("debug, agent_infos", agent_infos)
policy_time += time.time() - t
t = time.time()
next_obses, rewards, dones, env_infos = self.vec_env.step(actions, reset_args) # TODO: instead of receive obs from env, we'll receive it from the policy as a feed_dict
if contexts:
next_obses = np.concatenate([next_obses, contexts], axis = 1)
env_time += time.time() - t
t = time.time()
agent_infos = tensor_utils.split_tensor_dict_list(agent_infos)
env_infos = tensor_utils.split_tensor_dict_list(env_infos)
if env_infos is None:
env_infos = [dict() for _ in range(self.vec_env.num_envs)]
if agent_infos is None:
agent_infos = [dict() for _ in range(self.vec_env.num_envs)]
for idx, observation, action, reward, env_info, agent_info, done in zip(itertools.count(), obses, actions,
rewards, env_infos, agent_infos,
dones):
if running_paths[idx] is None:
running_paths[idx] = dict(
observations=[],
actions=[],
rewards=[],
env_infos=[],
agent_infos=[],
)
running_paths[idx]["observations"].append(observation)
running_paths[idx]["actions"].append(action)
running_paths[idx]["rewards"].append(reward)
running_paths[idx]["env_infos"].append(env_info)
running_paths[idx]["agent_infos"].append(agent_info)
if done:
paths[idx].append(dict(
observations=self.env_spec.observation_space.flatten_n(running_paths[idx]["observations"]),
actions=self.env_spec.action_space.flatten_n(running_paths[idx]["actions"]),
rewards=tensor_utils.stack_tensor_list(running_paths[idx]["rewards"]),
env_infos=tensor_utils.stack_tensor_dict_list(running_paths[idx]["env_infos"]),
agent_infos=tensor_utils.stack_tensor_dict_list(running_paths[idx]["agent_infos"]),
))
n_samples += len(running_paths[idx]["rewards"]) # TODO: let's also add the incomplete running_paths to paths
running_paths[idx] = None
path_nums[idx] += 1
process_time += time.time() - t
pbar.inc(len(obses))
obses = next_obses
# adding the incomplete paths
# for idx in range(self.vec_env.num_envs):
# if running_paths[idx] is not None:
# paths[idx].append(dict(
# observations=self.env_spec.observation_space.flatten_n(running_paths[idx]["observations"]),
# actions=self.env_spec.action_space.flatten_n(running_paths[idx]["actions"]),
# rewards=tensor_utils.stack_tensor_list(running_paths[idx]["rewards"]),
# env_infos=tensor_utils.stack_tensor_dict_list(running_paths[idx]["env_infos"]),
# agent_infos=tensor_utils.stack_tensor_dict_list(running_paths[idx]["agent_infos"]),
# ))
pbar.stop()
# logger.record_tabular(log_prefix + "PolicyExecTime", policy_time)
# logger.record_tabular(log_prefix + "EnvExecTime", env_time)
# logger.record_tabular(log_prefix + "ProcessExecTime", process_time)
if not return_dict:
flatten_list = lambda l: [item for sublist in l for item in sublist]
paths = flatten_list(paths.values())
#path_keys = flatten_list([[key]*len(paths[key]) for key in paths.keys()])
return paths
class VectorizedSampler(BaseSampler):
def __init__(self, algo, n_envs=None):
super(VectorizedSampler, self).__init__(algo)
self.n_envs = n_envs
def start_worker(self):
n_envs = self.n_envs
if n_envs is None:
n_envs = int(self.algo.batch_size / self.algo.max_path_length)
n_envs = max(1, min(n_envs, 100))
if getattr(self.algo.env, 'vectorized', False):
self.vec_env = self.algo.env.vec_env_executor(n_envs=n_envs, max_path_length=self.algo.max_path_length)
else:
envs = [pickle.loads(pickle.dumps(self.algo.env)) for _ in range(n_envs)]
self.vec_env = VecEnvExecutor(
envs=envs,
#env=pickle.loads(pickle.dumps(self.algo.env)),
#n = n_envs,
max_path_length=self.algo.max_path_length
)
self.env_spec = self.algo.env.spec
def shutdown_worker(self):
self.vec_env.terminate()
def obtain_samples(self, itr, reset_args=None, return_dict=False, log_prefix=''):
# reset_args: arguments to pass to the environments to reset
# return_dict: whether or not to return a dictionary or list form of paths
logger.log("Obtaining samples for iteration %d..." % itr)
#paths = []
paths = {}
for i in range(self.vec_env.num_envs):
paths[i] = []
# if the reset args are not list/numpy, we set the same args for each env
if reset_args is not None and (type(reset_args) != list and type(reset_args)!=np.ndarray):
reset_args = [reset_args]*self.vec_env.num_envs
n_samples = 0
obses = self.vec_env.reset(reset_args)
dones = np.asarray([True] * self.vec_env.num_envs)
running_paths = [None] * self.vec_env.num_envs
pbar = ProgBarCounter(self.algo.batch_size)
policy_time = 0
env_time = 0
process_time = 0
policy = self.algo.policy
import time
while n_samples < self.algo.batch_size:
t = time.time()
policy.reset(dones)
actions, agent_infos = policy.get_actions(obses)
policy_time += time.time() - t
t = time.time()
next_obses, rewards, dones, env_infos = self.vec_env.step(actions, reset_args)
env_time += time.time() - t
t = time.time()
agent_infos = tensor_utils.split_tensor_dict_list(agent_infos)
env_infos = tensor_utils.split_tensor_dict_list(env_infos)
if env_infos is None:
env_infos = [dict() for _ in range(self.vec_env.num_envs)]
if agent_infos is None:
agent_infos = [dict() for _ in range(self.vec_env.num_envs)]
for idx, observation, action, reward, env_info, agent_info, done in zip(itertools.count(), obses, actions,
rewards, env_infos, agent_infos,
dones):
if running_paths[idx] is None:
running_paths[idx] = dict(
observations=[],
actions=[],
rewards=[],
env_infos=[],
agent_infos=[],
)
running_paths[idx]["observations"].append(observation)
running_paths[idx]["actions"].append(action)
running_paths[idx]["rewards"].append(reward)
running_paths[idx]["env_infos"].append(env_info)
running_paths[idx]["agent_infos"].append(agent_info)
if done:
paths[idx].append(dict(
observations=self.env_spec.observation_space.flatten_n(running_paths[idx]["observations"]),
actions=self.env_spec.action_space.flatten_n(running_paths[idx]["actions"]),
rewards=tensor_utils.stack_tensor_list(running_paths[idx]["rewards"]),
env_infos=tensor_utils.stack_tensor_dict_list(running_paths[idx]["env_infos"]),
agent_infos=tensor_utils.stack_tensor_dict_list(running_paths[idx]["agent_infos"]),
))
n_samples += len(running_paths[idx]["rewards"])
running_paths[idx] = None
process_time += time.time() - t
pbar.inc(len(obses))
obses = next_obses
pbar.stop()
logger.record_tabular(log_prefix+"PolicyExecTime", policy_time)
logger.record_tabular(log_prefix+"EnvExecTime", env_time)
logger.record_tabular(log_prefix+"ProcessExecTime", process_time)
if not return_dict:
flatten_list = lambda l: [item for sublist in l for item in sublist]
paths = flatten_list(paths.values())
#path_keys = flatten_list([[key]*len(paths[key]) for key in paths.keys()])
return paths
class VectorizedSampler(BaseSampler):
def __init__(self, algo, n_envs=None):
super(VectorizedSampler, self).__init__(algo)
self.n_envs = n_envs
def start_worker(self, include_joint_coords=False):
n_envs = self.n_envs
if n_envs is None:
n_envs = int(self.algo.batch_size / self.algo.max_path_length)
n_envs = max(1, min(n_envs, 100))
if getattr(self.algo.env, 'vectorized', False):
self.vec_env = self.algo.env.vec_env_executor(
n_envs=n_envs,
max_path_length=self.algo.max_path_length,
include_joint_coords=include_joint_coords)
else:
envs = [
pickle.loads(pickle.dumps(self.algo.env))
for _ in range(n_envs)
]
self.vec_env = VecEnvExecutor(
envs=envs,
max_path_length=self.algo.max_path_length,
include_joint_coords=include_joint_coords)
self.env_spec = self.algo.env.spec
def shutdown_worker(self):
self.vec_env.terminate()
def obtain_samples_for_visualization(self, include_joint_coords=False):
tf_env = self.algo.env
if hasattr(tf_env.wrapped_env, "stats_recorder"):
setattr(tf_env.wrapped_env.stats_recorder, "done", None)
import builtins
builtins.visualize = True
print("\nAbout to start video...")
obs_dim = self.env_spec.observation_space.shape[0]
obs = tf_env.reset()
obs = self._add_joint_coords_to_obs(obs, include_joint_coords)
horizon = 1000
for horizon_num in range(1, horizon + 1):
# action, _ = self.algo.policy.get_action(obs[:obs_dim])
action, _ = self.algo.policy.get_action(obs)
next_obs, reward, done, _info = tf_env.step(action, use_states=obs)
obs = self._add_joint_coords_to_obs(next_obs, include_joint_coords)
if done or horizon_num == horizon:
break
builtins.visualize = False
def obtain_samples(self, itr, include_joint_coords=False):
logger.log("Obtaining samples for iteration %d..." % itr)
paths = []
n_samples = 0
obses = self.vec_env.reset()
obses = self._add_joint_coords_to_obses(obses, include_joint_coords)
obs_dim = self.env_spec.observation_space.shape[0]
dones = np.asarray([True] * self.vec_env.num_envs)
running_paths = [None] * self.vec_env.num_envs
pbar = ProgBarCounter(self.algo.batch_size)
policy_time = 0
env_time = 0
process_time = 0
policy = self.algo.policy
import time
while n_samples < self.algo.batch_size:
t = time.time()
policy.reset(dones)
actions, agent_infos = policy.get_actions(obses)
policy_time += time.time() - t
t = time.time()
next_obses, rewards, dones, env_infos = self.vec_env.step(
actions, use_states=obses)
env_time += time.time() - t
t = time.time()
agent_infos = tensor_utils.split_tensor_dict_list(agent_infos)
env_infos = tensor_utils.split_tensor_dict_list(env_infos)
if env_infos is None:
env_infos = [dict() for _ in range(self.vec_env.num_envs)]
if agent_infos is None:
agent_infos = [dict() for _ in range(self.vec_env.num_envs)]
for idx, observation, action, reward, env_info, agent_info, done in zip(
itertools.count(), obses, actions, rewards, env_infos,
agent_infos, dones):
if running_paths[idx] is None:
running_paths[idx] = dict(
observations=[],
actions=[],
rewards=[],
env_infos=[],
agent_infos=[],
)
running_paths[idx]["observations"].append(observation)
running_paths[idx]["actions"].append(action)
running_paths[idx]["rewards"].append(reward)
running_paths[idx]["env_infos"].append(env_info)
running_paths[idx]["agent_infos"].append(agent_info)
if done:
paths.append(
dict(
observations=self.env_spec.observation_space.
flatten_n(running_paths[idx]["observations"]),
actions=self.env_spec.action_space.flatten_n(
running_paths[idx]["actions"]),
rewards=tensor_utils.stack_tensor_list(
running_paths[idx]["rewards"]),
env_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["env_infos"]),
agent_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["agent_infos"]),
))
n_samples += len(running_paths[idx]["rewards"])
running_paths[idx] = None
process_time += time.time() - t
pbar.inc(len(obses))
obses = self._add_joint_coords_to_obses(next_obses,
include_joint_coords)
pbar.stop()
logger.record_tabular("PolicyExecTime", policy_time)
logger.record_tabular("EnvExecTime", env_time)
logger.record_tabular("ProcessExecTime", process_time)
return paths, n_samples
def _add_joint_coords_to_obses(self, obses, include_joint_coords):
if include_joint_coords:
try:
inner_env = self._get_inner_env()
extended_obses = []
for obs in obses:
extended_obses.append(
self._add_joint_coords_to_obs(obs,
include_joint_coords,
inner_env))
return np.array(extended_obses)
except AttributeError:
inner_envs = self._get_inner_envs()
extended_obses = []
for obs_i in range(len(obses)):
extended_obses.append(
self._add_joint_coords_to_obs(obses[obs_i],
include_joint_coords,
inner_envs[obs_i]))
return np.array(extended_obses)
return obses
def _add_joint_coords_to_obs(self,
obs,
include_joint_coords,
inner_env=None):
if include_joint_coords:
if not inner_env:
inner_env = self._get_inner_env()
if hasattr(inner_env, "get_geom_xpos"):
return np.append(obs, inner_env.get_geom_xpos().flatten())
else:
return np.append(obs, inner_env.env.get_geom_xpos().flatten())
return obs
def _get_inner_env(self):
env = self.vec_env.vec_env
while hasattr(env, "env"):
env = env.env
if hasattr(env.wrapped_env, '_wrapped_env'):
return env.wrapped_env._wrapped_env
else:
return env.wrapped_env.env.unwrapped
def _get_inner_envs(self):
inner_envs = []
for env in self.vec_env.envs:
while hasattr(env, "env"):
env = env.env
if hasattr(env.wrapped_env, '_wrapped_env'):
inner_envs.append(env.wrapped_env._wrapped_env)
else:
inner_envs.append(env.wrapped_env.env.unwrapped)
return inner_envs
class VectorizedSampler(BaseSampler):
def __init__(self, algo, n_envs=None):
super(VectorizedSampler, self).__init__(algo)
self.n_envs = n_envs
def start_worker(self):
n_envs = self.n_envs
if n_envs is None:
n_envs = int(self.algo.batch_size / self.algo.max_path_length)
n_envs = max(1, min(n_envs, 100))
if getattr(self.algo.env, 'vectorized', False):
self.vec_env = self.algo.env.vec_env_executor(
n_envs=n_envs, max_path_length=self.algo.max_path_length)
else:
envs = [
pickle.loads(pickle.dumps(self.algo.env))
for _ in range(n_envs)
]
self.vec_env = VecEnvExecutor(
envs=envs,
#env=pickle.loads(pickle.dumps(self.algo.env)),
#n = n_envs,
max_path_length=self.algo.max_path_length)
self.env_spec = self.algo.env.spec
def shutdown_worker(self):
self.vec_env.terminate()
def get_random_action(self, observation, idx=None):
# this function takes a numpy array observations and outputs randomly sampled actions.
# idx: index corresponding to the task/updated policy.
#flat_obs = self.env_spec.observation_space.flatten(observation)
#f_dist = self._cur_f_dist
mean = np.random.uniform(low=-1.0,
high=1.0,
size=[self.n_envs, self.algo.a_size])
#mean, log_std = [x[0] for x in f_dist([flat_obs])]
action = mean
return action, dict(mean=mean)
def get_MPC_action(self, obs):
action_list = []
for i in range(self.n_envs):
action, _ = self.algo.policy.get_action(obs[i], i)
action_list.append(action)
return np.asarray(action_list).reshape(
[-1, self.algo.env.action_space.shape[0]]), dict(means=action_list)
def obtain_samples(self,
itr,
init_state=None,
reset_args=None,
return_dict=False,
log_prefix=''):
# reset_args: arguments to pass to the environments to reset
# return_dict: whether or not to return a dictionary or list form of paths
logger.log("Obtaining samples for iteration %d..." % itr)
#paths = []
paths = {}
for i in range(self.vec_env.num_envs):
paths[i] = []
# if the reset args are not list/numpy, we set the same args for each env
if reset_args is not None and (type(reset_args) != list
and type(reset_args) != np.ndarray):
reset_args = [reset_args] * self.vec_env.num_envs
if init_state is not None:
init_state = [init_state] * self.vec_env.num_envs
n_samples = 0
obses = self.vec_env.reset(init_state, reset_args)
dones = np.asarray([True] * self.vec_env.num_envs)
running_paths = [None] * self.vec_env.num_envs
else:
n_samples = 0
obses = self.vec_env.reset(reset_args)
dones = np.asarray([True] * self.vec_env.num_envs)
running_paths = [None] * self.vec_env.num_envs
pbar = ProgBarCounter(self.algo.batch_size)
policy_time = 0
env_time = 0
process_time = 0
#policy = self.algo.policy
import time
while n_samples < self.algo.max_path_length:
t = time.time()
#self.env_spec.reset(reset_args = reset_args)
#policy.reset(dones)
actions, agent_infos = self.get_MPC_action(obses)
policy_time += time.time() - t
t = time.time()
next_obses, rewards, dones, env_infos = self.vec_env.step(
actions, reset_args)
env_time += time.time() - t
t = time.time()
agent_infos = tensor_utils.split_tensor_dict_list(agent_infos)
env_infos = tensor_utils.split_tensor_dict_list(env_infos)
if env_infos is None:
env_infos = [dict() for _ in range(self.vec_env.num_envs)]
if agent_infos is None:
agent_infos = [dict() for _ in range(self.vec_env.num_envs)]
for idx, observation, action, reward, env_info, agent_info, done in zip(
itertools.count(), obses, actions, rewards, env_infos,
agent_infos, dones):
if running_paths[idx] is None:
running_paths[idx] = dict(
observations=[],
actions=[],
rewards=[],
env_infos=[],
agent_infos=[],
)
running_paths[idx]["observations"].append(observation)
running_paths[idx]["actions"].append(action)
running_paths[idx]["rewards"].append(reward)
running_paths[idx]["env_infos"].append(env_info)
running_paths[idx]["agent_infos"].append(agent_info)
if done:
paths[idx].append(
dict(
observations=running_paths[idx]["observations"],
actions=running_paths[idx]["actions"],
rewards=tensor_utils.stack_tensor_list(
running_paths[idx]["rewards"]),
env_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["env_infos"]),
agent_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["agent_infos"]),
))
n_samples += len(running_paths[idx]["rewards"])
running_paths[idx] = None
process_time += time.time() - t
pbar.inc(len(obses))
obses = next_obses
pbar.stop()
logger.record_tabular(log_prefix + "PolicyExecTime", policy_time)
logger.record_tabular(log_prefix + "EnvExecTime", env_time)
logger.record_tabular(log_prefix + "ProcessExecTime", process_time)
if not return_dict:
flatten_list = lambda l: [
item for sublist in l for item in sublist
]
paths = flatten_list(paths.values())
#path_keys = flatten_list([[key]*len(paths[key]) for key in paths.keys()])
return paths
class VectorizedSampler(Sampler):
def __init__(self, env, policy, n_envs, vec_env=None):
self.env = env
self.policy = policy
self.n_envs = n_envs
self.vec_env = vec_env
self.env_spec = env.spec
def start_worker(self):
if self.vec_env is None:
n_envs = self.n_envs
if getattr(self.env, 'vectorized', False):
self.vec_env = self.env.vec_env_executor(n_envs=n_envs)
else:
envs = [
pickle.loads(pickle.dumps(self.env)) for _ in range(n_envs)
]
self.vec_env = VecEnvExecutor(envs=envs, )
def shutdown_worker(self):
self.vec_env.terminate()
def obtain_samples(self,
itr,
max_path_length,
batch_size,
max_n_trajs=None):
logger.log("Obtaining samples for iteration %d..." % itr)
paths = []
n_samples = 0
dones = np.asarray([True] * self.vec_env.n_envs)
obses = self.vec_env.reset(dones)
running_paths = [None] * self.vec_env.n_envs
pbar = ProgBarCounter(batch_size)
policy_time = 0
env_time = 0
process_time = 0
policy = self.policy
import time
while n_samples < batch_size:
t = time.time()
if hasattr(self.vec_env, "handle_policy_reset"):
self.vec_env.handle_policy_reset(policy, dones)
else:
policy.reset(dones)
actions, agent_infos = policy.get_actions(obses)
policy_time += time.time() - t
t = time.time()
next_obses, rewards, dones, env_infos = self.vec_env.step(
actions, max_path_length=max_path_length)
if np.any(dones):
new_obses = self.vec_env.reset(dones)
reset_idx = 0
for idx, done in enumerate(dones):
if done:
next_obses[idx] = new_obses[reset_idx]
reset_idx += 1
env_time += time.time() - t
t = time.time()
agent_infos = tensor_utils.split_tensor_dict_list(agent_infos)
env_infos = tensor_utils.split_tensor_dict_list(env_infos)
if env_infos is None:
env_infos = [dict() for _ in range(self.vec_env.n_envs)]
if agent_infos is None:
agent_infos = [dict() for _ in range(self.vec_env.n_envs)]
for idx, observation, action, reward, env_info, agent_info, done in zip(
itertools.count(), obses, actions, rewards, env_infos,
agent_infos, dones):
if running_paths[idx] is None:
running_paths[idx] = dict(
observations=[],
actions=[],
rewards=[],
env_infos=[],
agent_infos=[],
)
running_paths[idx]["observations"].append(observation)
running_paths[idx]["actions"].append(action)
running_paths[idx]["rewards"].append(reward)
running_paths[idx]["env_infos"].append(env_info)
running_paths[idx]["agent_infos"].append(agent_info)
if done:
paths.append(
dict(
observations=self.env_spec.observation_space.
flatten_n(running_paths[idx]["observations"]),
actions=self.env_spec.action_space.flatten_n(
running_paths[idx]["actions"]),
rewards=tensor_utils.stack_tensor_list(
running_paths[idx]["rewards"]),
env_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["env_infos"]),
agent_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["agent_infos"]),
))
n_samples += len(running_paths[idx]["rewards"])
running_paths[idx] = None
if max_n_trajs is not None and len(paths) >= max_n_trajs:
break
process_time += time.time() - t
pbar.inc(len(obses))
obses = next_obses
pbar.stop()
logger.record_tabular("PolicyExecTime", policy_time)
logger.record_tabular("EnvExecTime", env_time)
logger.record_tabular("ProcessExecTime", process_time)
return paths
class RandomVectorizedSampler(RandomBaseSampler, VectorizedSampler):
def __init__(self, algo, n_envs=None):
self.algo = algo
VectorizedSampler.__init__(self, algo, n_envs=n_envs)
def start_worker(self):
n_envs = self.n_envs
if n_envs is None:
n_envs = int(self.algo.batch_size / self.algo.max_path_length)
n_envs = max(1, min(n_envs, 100))
if getattr(self.algo.env, 'vectorized', False):
self.vec_env = self.algo.env.vec_env_executor(
n_envs=n_envs, max_path_length=self.algo.max_path_length)
else:
envs = [
pickle.loads(pickle.dumps(self.algo.env))
for _ in range(n_envs)
]
self.vec_env = VecEnvExecutor(
envs=envs, max_path_length=self.algo.max_path_length)
self.env_spec = self.algo.env.spec
def shutdown_worker(self):
self.vec_env.terminate()
@overrides
def obtain_samples(self,
itr,
num_samples=None,
log=True,
log_prefix='RandomSampler-'):
if num_samples is None:
num_samples = self.algo.batch_size
paths = []
n_samples_collected = 0
obses = self.vec_env.reset()
dones = np.asarray([True] * self.vec_env.num_envs)
running_paths = [None] * self.vec_env.num_envs
pbar = ProgBarCounter(num_samples)
env_time = 0
process_time = 0
policy = self.algo.policy
import time
while n_samples_collected < num_samples:
# random actions
t = time.time()
actions = np.stack([
self.vec_env.action_space.sample() for _ in range(len(obses))
],
axis=0)
policy_time = time.time() - t
agent_infos = {}
t = time.time()
next_obses, rewards, dones, env_infos = self.vec_env.step(actions)
env_time += time.time() - t
t = time.time()
agent_infos = tensor_utils.split_tensor_dict_list(agent_infos)
env_infos = tensor_utils.split_tensor_dict_list(env_infos)
if env_infos is None:
env_infos = [dict() for _ in range(self.vec_env.num_envs)]
if agent_infos is None:
agent_infos = [dict() for _ in range(self.vec_env.num_envs)]
for idx, observation, action, reward, env_info, agent_info, done in zip(
itertools.count(), obses, actions, rewards, env_infos,
agent_infos, dones):
if running_paths[idx] is None:
running_paths[idx] = dict(
observations=[],
actions=[],
rewards=[],
env_infos=[],
agent_infos=[],
)
running_paths[idx]["observations"].append(observation)
running_paths[idx]["actions"].append(action)
running_paths[idx]["rewards"].append(reward)
running_paths[idx]["env_infos"].append(env_info)
running_paths[idx]["agent_infos"].append(agent_info)
if done:
paths.append(
dict(
observations=self.env_spec.observation_space.
flatten_n(running_paths[idx]["observations"]),
actions=self.env_spec.action_space.flatten_n(
running_paths[idx]["actions"]),
rewards=tensor_utils.stack_tensor_list(
running_paths[idx]["rewards"]),
env_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["env_infos"]),
agent_infos=tensor_utils.stack_tensor_dict_list(
running_paths[idx]["agent_infos"]),
))
n_samples_collected += len(running_paths[idx]["rewards"])
running_paths[idx] = None
process_time += time.time() - t
pbar.inc(len(obses))
obses = next_obses
pbar.stop()
if log:
logger.record_tabular(log_prefix + "PolicyExecTime", policy_time)
logger.record_tabular(log_prefix + "EnvExecTime", env_time)
logger.record_tabular(log_prefix + "ProcessExecTime", process_time)
return paths