def __init__(self, env, max_obs_dim=None): super().__init__(env) self._max_obs_dim = max_obs_dim action_space = akro.from_gym(self.env.action_space) observation_space = self._create_rl2_obs_space(env) self._spec = EnvSpec(action_space=action_space, observation_space=observation_space)
def __init__(self, goal_reward=10, actuation_cost_coeff=30, distance_cost_coeff=1, init_sigma=0.1): super().__init__() Serializable.quick_init(self, locals()) self.dynamics = PointDynamics(dim=2, sigma=0) self.init_mu = np.zeros(2, dtype=np.float32) self.init_sigma = init_sigma self.goal_positions = np.array([[5, 0], [-5, 0], [0, 5], [0, -5]], dtype=np.float32) self.goal_threshold = 1. self.goal_reward = goal_reward self.action_cost_coeff = actuation_cost_coeff self.distance_cost_coeff = distance_cost_coeff self.xlim = (-7, 7) self.ylim = (-7, 7) self.vel_bound = 1. self.reset() self.observation = None self._ax = None self._env_lines = list() self.fixed_plots = None self.dynamic_plots = [] self.__spec = EnvSpec(action_space=self.action_space, observation_space=self.observation_space)
def test_pickleable(self): env_spec = EnvSpec(akro.Box(-1, 1, (1, )), akro.Box(-2, 2, (2, )), 500) round_trip = pickle.loads(pickle.dumps(env_spec)) assert round_trip assert round_trip.action_space == env_spec.action_space assert round_trip.observation_space == env_spec.observation_space assert round_trip.max_episode_length == env_spec.max_episode_length
def __init__(self, env=None, env_name='', is_image=False): """Returns a Garage wrapper class for bullet-based gym.Env. Args: env (gym.wrappers.time_limit): A gym.wrappers.time_limit.TimeLimit object wrapping a gym.Env created via gym.make(). env_name (str): If the env_name is speficied, a gym environment with that name will be created. If such an environment does not exist, a `gym.error` is thrown. is_image (bool): True if observations contain pixel values, false otherwise. Setting this to true converts a gym.Spaces.Box obs space to an akro.Image and normalizes pixel values. """ if not env: # 'RacecarZedBulletEnv-v0' environment enables rendering by # default, while pybullet allows only one GUI connection at a time. # Setting renders to False avoids potential error when multiple # of these envs are tested at the same time. if env_name == 'RacecarZedBulletEnv-v0': env = gym.make(env_name, renders=False) else: env = gym.make(env_name) # Needed for deserialization self._env = env self._env_name = env_name super().__init__(env) self.action_space = akro.from_gym(self.env.action_space) self.observation_space = akro.from_gym(self.env.observation_space, is_image=is_image) self._spec = EnvSpec(action_space=self.action_space, observation_space=self.observation_space)
def __init__(self, env=None, env_name='', is_image=False): """Initializes a GarageEnv. Args: env (gym.wrappers.time_limit): A gym.wrappers.time_limit.TimeLimit object wrapping a gym.Env created via gym.make(). env_name (str): If the env_name is speficied, a gym environment with that name will be created. If such an environment does not exist, a `gym.error` is thrown. is_image (bool): True if observations contain pixel values, false otherwise. Setting this to true converts a gym.Spaces.Box obs space to an akro.Image and normalizes pixel values. """ # Needed for deserialization self._env_name = env_name self._env = env if env_name: super().__init__(gym.make(env_name)) else: super().__init__(env) self.action_space = akro.from_gym(self.env.action_space) self.observation_space = akro.from_gym(self.env.observation_space, is_image=is_image) self._spec = EnvSpec(action_space=self.action_space, observation_space=self.observation_space)
def spec(self): """ Returns an EnvSpec. Returns: spec (garage.envs.EnvSpec) """ return EnvSpec(observation_space=self.observation_space, action_space=self.action_space)
def spec(self): r"""Returns a garage environment specification. Returns ------- :py:class:`garage.envs.env_spec.EnvSpec` A garage environment specification. """ return EnvSpec( observation_space=self.observation_space, action_space=self.action_space)
def __init__(self, env, task_index, n_total_tasks): assert 0 <= task_index < n_total_tasks super().__init__(env) self._task_index = task_index self._n_total_tasks = n_total_tasks env_lb = self.env.observation_space.low env_ub = self.env.observation_space.high one_hot_ub = np.ones(self._n_total_tasks) one_hot_lb = np.zeros(self._n_total_tasks) self.observation_space = akro.Box(np.concatenate([env_lb, one_hot_lb]), np.concatenate([env_ub, one_hot_ub])) self.__spec = EnvSpec(action_space=self.action_space, observation_space=self.observation_space)
def __init__(self, env=None, env_name=''): # Needed for deserialization self._env_name = env_name self._env = env if env_name: super().__init__(gym.make(env_name)) else: super().__init__(env) self.action_space = akro.from_gym(self.env.action_space) self.observation_space = akro.from_gym(self.env.observation_space) self.__spec = EnvSpec(action_space=self.action_space, observation_space=self.observation_space)
def __init__(self, env=None, env_name=''): if env_name: super().__init__(gym.make(env_name)) else: super().__init__(env) self.action_space = akro.from_gym(self.env.action_space) self.observation_space = akro.from_gym(self.env.observation_space) if self.spec: self.spec.action_space = self.action_space self.spec.observation_space = self.observation_space else: self.spec = EnvSpec(action_space=self.action_space, observation_space=self.observation_space) Serializable.quick_init(self, locals())
def __init__(self, env=None, env_name='', is_image=False): # Needed for deserialization self._env_name = env_name self._env = env if env_name: print(env_name) # @TODO fix this line blowing things up super().__init__(gym.make(env_name)) else: super().__init__(env) self.action_space = akro.from_gym(self.env.action_space) self.observation_space = akro.from_gym(self.env.observation_space, is_image=is_image) self.__spec = EnvSpec(action_space=self.action_space, observation_space=self.observation_space)
def setup_method(self): """Setup for all test methods.""" self.latent_dim = 5 self.env_spec = TfEnv(DummyBoxEnv()) latent_space = akro.Box(low=-1, high=1, shape=(self.latent_dim, ), dtype=np.float32) # add latent space to observation space to create a new space augmented_obs_space = akro.Tuple( (self.env_spec.observation_space, latent_space)) augmented_env_spec = EnvSpec(augmented_obs_space, self.env_spec.action_space) self.obs_dim = int(np.prod(self.env_spec.observation_space.shape)) self.action_dim = int(np.prod(self.env_spec.action_space.shape)) reward_dim = 1 self.encoder_input_dim = self.obs_dim + self.action_dim + reward_dim encoder_output_dim = self.latent_dim * 2 encoder_hidden_sizes = (3, 2, encoder_output_dim) context_encoder = MLPEncoder(input_dim=self.encoder_input_dim, output_dim=encoder_output_dim, hidden_nonlinearity=None, hidden_sizes=encoder_hidden_sizes, hidden_w_init=nn.init.ones_, output_w_init=nn.init.ones_) context_policy = TanhGaussianMLPPolicy(env_spec=augmented_env_spec, hidden_sizes=(3, 5, 7), hidden_nonlinearity=F.relu, output_nonlinearity=None) self.module = ContextConditionedPolicy(latent_dim=self.latent_dim, context_encoder=context_encoder, policy=context_policy, use_information_bottleneck=True, use_next_obs=False)
def __init__(self, env=None, env_name='', is_image=False): # Needed for deserialization self._env_name = env_name self._env = env if env_name: super().__init__(gym.make(env_name)) else: super().__init__(env) if isinstance(self.env.action_space, Box): self.action_space = akro.Box(low=self.env.action_space.low, high=self.env.action_space.high) self.observation_space = akro.Image( shape=self.env.observation_space.shape) else: self.action_space = akro.from_gym(self.env.action_space) self.observation_space = akro.from_gym(self.env.observation_space, is_image=is_image) self.__spec = EnvSpec(action_space=self.action_space, observation_space=self.observation_space)
def test_init_std(self): sess = tf.Session() sess.__enter__() task_space = Box( np.zeros(8, dtype=np.float32), np.ones(8, dtype=np.float32)) latent_space = Box( np.zeros(4, dtype=np.float32), np.ones(4, dtype=np.float32)) embed_spec = EmbeddingSpec(task_space, latent_space) embedding = GaussianMLPEmbedding( name="embedding", embedding_spec=embed_spec, hidden_sizes=(20, 20), std_share_network=True, init_std=1.0, max_std=2.0, ) std_parameterization = "exp" observation_space = Box(np.full(100, 0.0), np.full(100, 100.0)) action_space = Box(np.full(10, 0.0), np.full(10, 10.0)) env_spec = EnvSpec(observation_space, action_space) policy = GaussianMLPMultitaskPolicy( name="policy", env_spec=env_spec, task_space=task_space, embedding=embedding, hidden_sizes=(200, 100), std_share_network=True, # max_std=10.0, init_std=6.0, std_parameterization=std_parameterization, ) sess.run(tf.global_variables_initializer()) z = latent_space.sample() # z = np.ones_like(latent_space.low) print("|z| = {}".format(np.linalg.norm(z))) # z = z / np.linalg.norm(z) o = observation_space.sample() # o = np.ones_like(observation_space.low) print("|o| = {}".format(np.linalg.norm(o))) # o = o / np.linalg.norm(o) a, info = policy.get_action_from_latent(z, o) log_stds = info["log_std"] if std_parameterization == "exp": stds = np.exp(log_stds) elif std_parameterization == "softplus": stds = np.log(1. + np.exp(log_stds)) else: raise NotImplementedError print("log_stds = {}".format(log_stds)) print("stds = {}".format(stds)) print("mean(stds) = {}".format(np.mean(stds))) print("std(stds) = {}".format(np.std(stds))) assert np.allclose(stds, 1.0), "stds: {}".format(stds)
def run_task(v): v = SimpleNamespace(**v) task_names = sorted(v.tasks.keys()) task_args = [v.tasks[t]['args'] for t in task_names] task_kwargs = [v.tasks[t]['kwargs'] for t in task_names] with TaskEmbeddingRunner() as runner: # Environment env = TfEnv( normalize( MultiTaskEnv( task_env_cls=FlatTorqueReacher, task_args=task_args, task_kwargs=task_kwargs))) # Latent space and embedding specs # TODO(gh/10): this should probably be done in Embedding or Algo latent_lb = np.zeros(v.latent_length, ) latent_ub = np.ones(v.latent_length, ) latent_space = Box(latent_lb, latent_ub) # trajectory space is (TRAJ_ENC_WINDOW, act_obs) where act_obs is a stacked # vector of flattened actions and observations act_lb, act_ub = env.action_space.bounds act_lb_flat = env.action_space.flatten(act_lb) act_ub_flat = env.action_space.flatten(act_ub) obs_lb, obs_ub = env.observation_space.bounds obs_lb_flat = env.observation_space.flatten(obs_lb) obs_ub_flat = env.observation_space.flatten(obs_ub) # act_obs_lb = np.concatenate([act_lb_flat, obs_lb_flat]) # act_obs_ub = np.concatenate([act_ub_flat, obs_ub_flat]) act_obs_lb = obs_lb_flat act_obs_ub = obs_ub_flat # act_obs_lb = act_lb_flat # act_obs_ub = act_ub_flat traj_lb = np.stack([act_obs_lb] * v.inference_window) traj_ub = np.stack([act_obs_ub] * v.inference_window) traj_space = Box(traj_lb, traj_ub) task_embed_spec = EmbeddingSpec(env.task_space, latent_space) traj_embed_spec = EmbeddingSpec(traj_space, latent_space) task_obs_space = concat_spaces(env.task_space, env.observation_space) env_spec_embed = EnvSpec(task_obs_space, env.action_space) # TODO(): rename to inference_network traj_embedding = GaussianMLPEmbedding( name="inference", embedding_spec=traj_embed_spec, hidden_sizes=(64, 64), std_share_network=True, init_std=1.0, ) # Embeddings task_embedding = GaussianMLPEmbedding( name="embedding", embedding_spec=task_embed_spec, hidden_sizes=(64, 64), std_share_network=True, init_std=v.embedding_init_std, # 1.0 max_std=v.embedding_max_std, # 2.0 # std_parameterization="softplus", ) # Multitask policy policy = GaussianMLPMultitaskPolicy( name="policy", env_spec=env.spec, task_space=env.task_space, embedding=task_embedding, hidden_sizes=(64, 32), std_share_network=True, init_std=v.policy_init_std, max_std=v.policy_max_std, # std_parameterization="softplus", ) # baseline = MultiTaskLinearFeatureBaseline(env_spec=env_spec_embed) extra = v.latent_length + len(v.tasks) baseline = MultiTaskGaussianMLPBaseline( env_spec=env.spec, extra_dims=extra) algo = PPOTaskEmbedding( env=env, policy=policy, baseline=baseline, inference=traj_embedding, batch_size=v.batch_size, # 4096 max_path_length=v.max_path_length, discount=0.99, lr_clip_range=0.2, policy_ent_coeff=v.policy_ent_coeff, embedding_ent_coeff=v.embedding_ent_coeff, inference_ce_coeff=v.inference_ce_coeff, #optimizer_args=dict(max_grad_norm=0.5) ) runner.setup(algo, env, batch_size=v.batch_size, max_path_length=v.max_path_length) runner.train(n_epochs=1000, plot=False)
def run_task(v): v = SimpleNamespace(**v) task_names = sorted(v.tasks.keys()) task_args = [v.tasks[t]['args'] for t in task_names] task_kwargs = [v.tasks[t]['kwargs'] for t in task_names] with TaskEmbeddingRunner() as runner: # Environment env = TfEnv( MultiTaskEnv(task_env_cls=SimpleReacherEnv, task_args=task_args, task_kwargs=task_kwargs)) # Latent space and embedding specs # TODO(gh/10): this should probably be done in Embedding or Algo latent_lb = np.zeros(v.latent_length, ) latent_ub = np.ones(v.latent_length, ) latent_space = Box(latent_lb, latent_ub) # trajectory space is (TRAJ_ENC_WINDOW, act_obs) where act_obs is a stacked # vector of flattened actions and observations act_lb, act_ub = env.action_space.bounds act_lb_flat = env.action_space.flatten(act_lb) act_ub_flat = env.action_space.flatten(act_ub) obs_lb, obs_ub = env.observation_space.bounds obs_lb_flat = env.observation_space.flatten(obs_lb) obs_ub_flat = env.observation_space.flatten(obs_ub) # act_obs_lb = np.concatenate([act_lb_flat, obs_lb_flat]) # act_obs_ub = np.concatenate([act_ub_flat, obs_ub_flat]) act_obs_lb = obs_lb_flat act_obs_ub = obs_ub_flat # act_obs_lb = act_lb_flat # act_obs_ub = act_ub_flat traj_lb = np.stack([act_obs_lb] * v.inference_window) traj_ub = np.stack([act_obs_ub] * v.inference_window) traj_space = Box(traj_lb, traj_ub) task_embed_spec = EmbeddingSpec(env.task_space, latent_space) traj_embed_spec = EmbeddingSpec(traj_space, latent_space) task_obs_space = concat_spaces(env.task_space, env.observation_space) env_spec_embed = EnvSpec(task_obs_space, env.action_space) # TODO(): rename to inference_network traj_embedding = GaussianMLPEmbedding( name="inference", embedding_spec=traj_embed_spec, hidden_sizes=(200, 100), # was the same size as policy in Karol's paper std_share_network=True, init_std=2.0, ) # Embeddings task_embedding = GaussianMLPEmbedding( name="embedding", embedding_spec=task_embed_spec, hidden_sizes=(200, 200), std_share_network=True, init_std=v.embedding_init_std, max_std=v.embedding_max_std, ) # Multitask policy policy = GaussianMLPMultitaskPolicy( name="policy", env_spec=env.spec, task_space=env.task_space, embedding=task_embedding, hidden_sizes=(200, 100), std_share_network=True, init_std=v.policy_init_std, ) extra = v.latent_length + len(v.tasks) baseline = MultiTaskGaussianMLPBaseline( env_spec=env.spec, extra_dims=extra, regressor_args=dict(hidden_sizes=(200, 100)), ) algo = PPOTaskEmbedding( env=env, policy=policy, baseline=baseline, inference=traj_embedding, max_path_length=v.max_path_length, n_itr=2000, discount=0.99, lr_clip_range=0.2, policy_ent_coeff=v.policy_ent_coeff, embedding_ent_coeff=v.embedding_ent_coeff, inference_ce_coeff=v.inference_ce_coeff, use_softplus_entropy=True, ) runner.setup(algo, env, batch_size=v.batch_size, max_path_length=v.max_path_length) runner.train(n_epochs=2000, plot=False)
data = [np.sin(matrices[i]) for i in range(100)] obs = [{'observations': [x], 'returns': [np.mean(x)]} for x in data] observations = np.concatenate([p['observations'] for p in obs]) returns = np.concatenate([p['returns'] for p in obs]) returns = returns.reshape((-1, 1)) paths = {'observations': [np.sin(matrices[i]) for i in range(100, 110)]} expected = [[np.mean(x)] for x in paths['observations']] return (obs, observations, returns), (paths, expected) test_env_spec = EnvSpec(observation_space=akro.Box(low=-1, high=1, shape=(10, 10, 3)), action_space=None) class TestGaussianCNNBaseline(TfGraphTestCase): @pytest.mark.large def test_fit_normalized(self): gcr = GaussianCNNBaseline(env_spec=test_env_spec, filters=((3, (3, 3)), (6, (3, 3))), strides=(1, 1), padding='SAME', hidden_sizes=(32, ), adaptive_std=False, use_trust_region=True)
def test_pickleable(self): env_spec = EnvSpec(Box(-1, 1, (1)), Box(-2, 2, (2))) round_trip = pickle.loads(pickle.dumps(env_spec)) assert round_trip assert round_trip.action_space == env_spec.action_space assert round_trip.observation_space == env_spec.observation_space
def run_task(v): v = SimpleNamespace(**v) task_names = sorted(v.tasks.keys()) task_args = [v.tasks[t]['args'] for t in task_names] task_kwargs = [v.tasks[t]['kwargs'] for t in task_names] # Environment env = TfEnv( normalize( MultiTaskEnv( task_env_cls=Point3dEnv, task_args=task_args, task_kwargs=task_kwargs))) # Latent space and embedding specs # TODO(gh/10): this should probably be done in Embedding or Algo latent_lb = np.zeros(v.latent_length, ) latent_ub = np.ones(v.latent_length, ) latent_space = Box(latent_lb, latent_ub) # trajectory space is (TRAJ_ENC_WINDOW, act_obs) where act_obs is a stacked # vector of flattened actions and observations act_lb, act_ub = env.action_space.bounds act_lb_flat = env.action_space.flatten(act_lb) act_ub_flat = env.action_space.flatten(act_ub) obs_lb, obs_ub = env.observation_space.bounds obs_lb_flat = env.observation_space.flatten(obs_lb) obs_ub_flat = env.observation_space.flatten(obs_ub) # act_obs_lb = np.concatenate([act_lb_flat, obs_lb_flat]) # act_obs_ub = np.concatenate([act_ub_flat, obs_ub_flat]) act_obs_lb = obs_lb_flat act_obs_ub = obs_ub_flat # act_obs_lb = act_lb_flat # act_obs_ub = act_ub_flat traj_lb = np.stack([act_obs_lb] * v.inference_window) traj_ub = np.stack([act_obs_ub] * v.inference_window) traj_space = Box(traj_lb, traj_ub) task_embed_spec = EmbeddingSpec(env.task_space, latent_space) traj_embed_spec = EmbeddingSpec(traj_space, latent_space) task_obs_space = concat_spaces(env.task_space, env.observation_space) env_spec_embed = EnvSpec(task_obs_space, env.action_space) # TODO(): rename to inference_network traj_embedding = GaussianMLPEmbedding( name="inference", embedding_spec=traj_embed_spec, hidden_sizes=(20, 20), std_share_network=True, init_std=1.0, ) # Embeddings task_embedding = GaussianMLPEmbedding( name="embedding", embedding_spec=task_embed_spec, hidden_sizes=(20, 20), std_share_network=True, init_std=1.0, max_std=2.0, # normalize=True, ) # Multitask policy policy = GaussianMLPMultitaskPolicy( name="policy", env_spec=env.spec, task_space=env.task_space, embedding=task_embedding, hidden_sizes=(20, 10), std_share_network=True, # max_std=6.0, init_std=6.0, ) # baseline = MultiTaskLinearFeatureBaseline(env_spec=env_spec_embed) extra = v.latent_length + len(v.tasks) baseline = MultiTaskGaussianMLPBaseline( env_spec=env.spec, extra_dims=extra) algo = PPOTaskEmbedding( env=env, policy=policy, baseline=baseline, inference=traj_embedding, batch_size=v.batch_size, # 4096 max_path_length=50, n_itr=500, discount=0.99, step_size=0.2, plot=True, policy_ent_coeff=v.policy_ent_coeff, embedding_ent_coeff=v.embedding_ent_coeff, inference_ce_coeff=v.inference_ce_coeff, num_tasks_held_out=1, ) algo.train()
def run_task(*_): with TaskEmbeddingRunner() as runner: # Environment env = TfEnv( MultiTaskEnv( task_env_cls=PointEnv, task_args=TASK_ARGS, task_kwargs=TASK_KWARGS)) # Latent space and embedding specs # TODO(gh/10): this should probably be done in Embedding or Algo latent_lb = np.zeros(LATENT_LENGTH, ) latent_ub = np.ones(LATENT_LENGTH, ) latent_space = Box(latent_lb, latent_ub) # trajectory space is (TRAJ_ENC_WINDOW, act_obs) where act_obs is a stacked # vector of flattened actions and observations act_lb, act_ub = env.action_space.bounds act_lb_flat = env.action_space.flatten(act_lb) act_ub_flat = env.action_space.flatten(act_ub) obs_lb, obs_ub = env.observation_space.bounds obs_lb_flat = env.observation_space.flatten(obs_lb) obs_ub_flat = env.observation_space.flatten(obs_ub) # act_obs_lb = np.concatenate([act_lb_flat, obs_lb_flat]) # act_obs_ub = np.concatenate([act_ub_flat, obs_ub_flat]) act_obs_lb = obs_lb_flat act_obs_ub = obs_ub_flat # act_obs_lb = act_lb_flat # act_obs_ub = act_ub_flat traj_lb = np.stack([act_obs_lb] * TRAJ_ENC_WINDOW) traj_ub = np.stack([act_obs_ub] * TRAJ_ENC_WINDOW) traj_space = Box(traj_lb, traj_ub) task_embed_spec = EmbeddingSpec(env.task_space, latent_space) traj_embed_spec = EmbeddingSpec(traj_space, latent_space) task_obs_space = concat_spaces(env.task_space, env.observation_space) env_spec_embed = EnvSpec(task_obs_space, env.action_space) # Embeddings task_embedding = GaussianMLPEmbedding( name="embedding", embedding_spec=task_embed_spec, hidden_sizes=(20, 20), std_share_network=True, init_std=3.0, # 2.0 ) # TODO(): rename to inference_network traj_embedding = GaussianMLPEmbedding( name="inference", embedding_spec=traj_embed_spec, hidden_sizes=(20, 10), # was the same size as policy in Karol's paper std_share_network=True, ) # Multitask policy policy = GaussianMLPMultitaskPolicy( name="policy", env_spec=env.spec, task_space=env.task_space, embedding=task_embedding, hidden_sizes=(20, 10), std_share_network=True, # Must be True for embedding learning init_std=6.0, # 4.5 6.0 ) baseline = MultiTaskLinearFeatureBaseline(env_spec=env_spec_embed) max_path_length = 50 algo = TRPOTaskEmbedding( env=env, policy=policy, baseline=baseline, inference=traj_embedding, max_path_length=max_path_length, discount=0.99, max_kl_step=0.2, policy_ent_coeff=1e-7, # 1e-7 embedding_ent_coeff=1e-3, # 1e-3 inference_ce_coeff=1e-7, # 1e-7 # kl_constraint=KLConstraint.SOFT, # optimizer_args=dict(max_penalty=1e9), ) runner.setup(algo, env, batch_size=20000, max_path_length=max_path_length) runner.train(n_epochs=1000, plot=False)
def test_methods(): """Test PEARLWorker methods.""" env_spec = TfEnv(DummyBoxEnv()) latent_dim = 5 latent_space = akro.Box(low=-1, high=1, shape=(latent_dim, ), dtype=np.float32) # add latent space to observation space to create a new space augmented_obs_space = akro.Tuple( (env_spec.observation_space, latent_space)) augmented_env_spec = EnvSpec(augmented_obs_space, env_spec.action_space) obs_dim = int(np.prod(env_spec.observation_space.shape)) action_dim = int(np.prod(env_spec.action_space.shape)) reward_dim = 1 encoder_input_dim = obs_dim + action_dim + reward_dim encoder_output_dim = latent_dim * 2 encoder_hidden_sizes = (3, 2, encoder_output_dim) context_encoder = MLPEncoder(input_dim=encoder_input_dim, output_dim=encoder_output_dim, hidden_nonlinearity=None, hidden_sizes=encoder_hidden_sizes, hidden_w_init=nn.init.ones_, output_w_init=nn.init.ones_) policy = TanhGaussianMLPPolicy(env_spec=augmented_env_spec, hidden_sizes=(3, 5, 7), hidden_nonlinearity=F.relu, output_nonlinearity=None) context_policy = ContextConditionedPolicy(latent_dim=latent_dim, context_encoder=context_encoder, policy=policy, use_information_bottleneck=True, use_next_obs=False) max_path_length = 20 worker1 = PEARLWorker(seed=1, max_path_length=max_path_length, worker_number=1) worker1.update_agent(context_policy) worker1.update_env(env_spec) rollouts = worker1.rollout() assert rollouts.observations.shape == (max_path_length, obs_dim) assert rollouts.actions.shape == (max_path_length, action_dim) assert rollouts.rewards.shape == (max_path_length, ) worker2 = PEARLWorker(seed=1, max_path_length=max_path_length, worker_number=1, deterministic=True, accum_context=True) worker2.update_agent(context_policy) worker2.update_env(env_spec) rollouts = worker2.rollout() assert context_policy.context.shape == (1, max_path_length, encoder_input_dim) assert rollouts.observations.shape == (max_path_length, obs_dim) assert rollouts.actions.shape == (max_path_length, action_dim) assert rollouts.rewards.shape == (max_path_length, )