update algos with storage

pyrobolearn/algos/explorer.py

@@ -137,26 +137,32 @@ class Explorer(object):
         obs = self.env.reset()
         print("\nExplorer - initial state: {}".format(obs))
 
+        # reset storage
+        self.storage.reset(init_observations=obs)
+
         # reset explorer
         self.explorer.reset()
 
         # run RL task for T steps
         for step in range(num_steps):
             # get action and corresponding distribution from policy
-            act, dist = self.explorer.act(obs, deterministic)
+            act, dist = self.explorer.act(obs, deterministic=deterministic)
 
             # perform one step in the environment
             next_obs, reward, done, info = self.env.step(act)
 
             # insert in storage
-            # print("\nExplorer:")
-            # print("1. Observation data: {}".format(obs.merged_torch_data))
-            # print("2. Action data: {}".format(act))
-            # print("3. Next observation data: {}".format(next_obs.merged_torch_data))
-            # print("4. Reward: {}".format(reward))
-            # print("5. \\pi(.|s): {}".format(dist))
-            # print("6. log \\pi(a|s): {}".format(dist.log_prob(act)))
-            self.storage.insert(obs.merged_torch_data, act.merged_torch_data, next_obs.data, reward, dist)
+            print("\nExplorer:")
+            print("1. Observation data: {}".format(obs))  # .merged_torch_data))
+            print("2. Action data: {}".format(act))
+            print("3. Next observation data: {}".format(next_obs))  # merged_torch_data))
+            print("4. Reward: {}".format(reward))
+            print("5. \\pi(.|s): {}".format(dist))
+            print("6. log \\pi(a|s): {}".format([d.log_prob(act) for d in dist]))
+
+            self.storage.insert(next_obs, act, reward, masks=done, distributions=dist)
+
+            raw_input('enter')
 
             obs = next_obs
             if done:

pyrobolearn/algos/rl_algo.py

@@ -319,10 +319,10 @@ class RLAlgo(object):  # Algo):
         history = {}
 
         # set the policy in training mode
-        self.policy.train(mode=True)
+        self.policy.train()
 
         # for each episode
-        for ep in range(num_episodes):
+        for episode in range(num_episodes):
 
             # for each rollout
             for rollout in range(num_rollouts):
@@ -337,7 +337,7 @@ class RLAlgo(object):  # Algo):
                 history.setdefault('loss', []).append(loss)
 
         # set the policy in test mode
-        self.policy.train(mode=False)
+        self.policy.eval()
 
         return history
 

pyrobolearn/storages/storage.py

@@ -735,10 +735,14 @@ class RolloutStorage(DictStorage):
         logger.debug('creating space for masks')
         self.create_new_entry('masks', shapes=1, num_steps=self.num_steps + 1)
 
+        # allocate space for action distribution
+        self.create_new_entry('distributions', shapes=[() for _ in action_shapes], num_steps=self.num_steps,
+                              dtype=object)
+
         # space for log probabilities on policy, distributions, scalar values from value functions,
         # recurrent hidden states, and others have to be allocated outside the class
 
-    def reset(self):
+    def reset(self, init_observations=None):
         """Reset the storage by copying the last value and setting it to the first value."""
         # for key, value in self.iteritems():
         #     if isinstance(value, list):
@@ -747,8 +751,14 @@ class RolloutStorage(DictStorage):
         #                 item[0].copy_(item[-1])
         #     elif isinstance(value, torch.Tensor) and len(value) == self.num_steps + 1:
         #             self[key][0].copy_(self[key][-1])
-        for observation in self.observations:
-            observation[0].copy_(observation[-1])
+        if init_observations is None:
+            for observation in self.observations:
+                observation[0].copy_(observation[-1])
+        else:
+            if not isinstance(init_observations, list):
+                init_observations = [init_observations]
+            for observation, value in zip(self.observations, init_observations):
+                observation[0].copy_(self._convert_to_tensor(value))
         self.masks[0].copy_(self.masks[-1])
         # self.recurrent_hidden_states[0].copy_(self.recurrent_hidden_states[-1])
 
@@ -804,7 +814,7 @@ class RolloutStorage(DictStorage):
             else:
                 set_tensor(self[key], step, values, copy=copy)
 
-    def insert(self, observations, actions, rewards, masks=None, update_step=True, **kwargs):
+    def insert(self, observations, actions, rewards, masks, distributions=None, update_step=True, **kwargs):
         # distributions, values=None):
         # recurrent_hidden_state, action_log_prob):
         """
@@ -815,6 +825,7 @@ class RolloutStorage(DictStorage):
             actions (torch.Tensor, list of torch.Tensor): (list of) action(s)
             rewards (float, int, torch.Tensor): reward value
             masks (float, int, torch.Tensor): masks. They are set to zeros after an episode has terminated.
+            distributions (torch.distributions.Distribution, None): action distribution.
             update_step (bool): if True, it will update the current time step. If False, the user needs to call
                 `step()` in order to update it.
             **kwargs (dict): dictionary containing other parameters to update in the storage. The other parameters
@@ -828,6 +839,8 @@ class RolloutStorage(DictStorage):
             observations = [observations]
         if not isinstance(actions, list):
             actions = [actions]
+        if not isinstance(distributions, list):
+            distributions = [distributions]
 
         # insert each observation / action
         for observation, storage in zip(observations, self.observations):
@@ -841,6 +854,10 @@ class RolloutStorage(DictStorage):
             masks = torch.tensor(1.)
         self.masks[self._step + 1].copy_(self._convert_to_tensor(masks))
 
+        # insert distributions
+        for distribution, storage in zip(distributions, self.distributions):
+            storage[self._step] = distribution
+
         # add other elements
         for key, value in kwargs:
             if key in self and key in self._shifts: