add metrics (ongoing), fix minor errors and refactor rl algos

pyrobolearn/actions/__init__.py

@@ -2,6 +2,9 @@
 # import action
 from .action import Action
 
+# import basic actions
+from .basic_actions import *
+
 # import robot actions
 from .robot_actions import *
 

pyrobolearn/actions/action.py

@@ -36,7 +36,7 @@ class Action(object):
 
     In our framework, the `Action` class is decoupled from the policy and environment rendering it more modular [1].
     Nevertheless, the `Action` class still acts as a bridge between the policy and environment. In addition to be
-    the output of a policy/controller, it can also be the input to some value estimators, dynamic models, reward
+    the output of a policy/controller, it can also be the input to some value approximators, dynamic models, reward
     functions, and so on.
 
     This class also describes the `action_space` which has initially been defined in `gym.Env` [2].

pyrobolearn/actions/basic_actions.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+"""Define basic actions
+
+This includes notably the fixed and functional actions.
+"""
+
+import numpy as np
+
+from pyrobolearn.actions import Action
+
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class FixedAction(Action):
+    r"""Fixed Action.
+
+    This is a dummy fixed action which always returns the value it was initialized with.
+    """
+
+    def __init__(self, value):
+        super(FixedAction, self).__init__(data=value)
+
+    def _write(self, data=None):
+        pass
+
+
+class FunctionalAction(Action):
+    r"""Functional Action.
+
+    This is an action which accepts a function which has to output the data.
+    """
+
+    def __init__(self, function, initial_data):
+        self.function = function
+        super(FunctionalAction, self).__init__(data=initial_data)
+
+    def _write(self, data=None):
+        self.data = self.function(data)
+

pyrobolearn/algos/evaluator.py

@@ -62,6 +62,11 @@ class Evaluator(object):
         """Return the storage unit."""
         return self.estimator.storage
 
+    @storage.setter
+    def storage(self, storage):
+        """Set the storage unit."""
+        self.estimator.storage = storage
+
     ###########
     # Methods #
     ###########

pyrobolearn/algos/explorer.py

@@ -139,7 +139,7 @@ class Explorer(object):
         observation = self.env.reset()
         if verbose:
             print("\n#### Starting the Exploration phase ####")
-            print("Explorer - initial state: {}".format(observation))
+            # print("Explorer - initial state: {}".format(observation))
 
         # reset storage
         self.storage.reset(init_states=observation, rollout_idx=rollout_idx)
@@ -155,14 +155,14 @@ class Explorer(object):
             # perform one step in the environment
             next_observation, reward, done, info = self.env.step(action)
 
-            if verbose:
-                print("\nExplorer:")
-                print("1. Observation data: {}".format(observation))
-                print("2. Action data: {}".format(action))
-                print("3. Next observation data: {}".format(next_observation))
-                print("4. Reward: {}".format(reward))
-                print("5. \\pi(.|s): {}".format(distribution))
-                print("6. log \\pi(a|s): {}".format([d.log_prob(action) for d in distribution]))
+            # if verbose:
+            #     print("\nExplorer:")
+            #     print("1. Observation data: {}".format(observation))
+            #     print("2. Action data: {}".format(action))
+            #     print("3. Next observation data: {}".format(next_observation))
+            #     print("4. Reward: {}".format(reward))
+            #     print("5. \\pi(.|s): {}".format(distribution))
+            #     print("6. log \\pi(a|s): {}".format([d.log_prob(action) for d in distribution]))
 
             # insert in storage
             self.storage.insert(observation, action, next_observation, reward, mask=(1-done),

pyrobolearn/algos/ppo.py

@@ -16,7 +16,7 @@ from pyrobolearn.exploration import ActionExploration
 
 from pyrobolearn.storages import RolloutStorage
 from pyrobolearn.samplers import BatchRandomSampler
-from pyrobolearn.returns import GAE
+from pyrobolearn.returns import GAE, PolicyEvaluator
 from pyrobolearn.losses import CLIPLoss, L2Loss, EntropyLoss
 from pyrobolearn.optimizers import Adam
 
@@ -171,8 +171,9 @@ class PPO(GradientRLAlgo):
         if not isinstance(actor_critic, ActorCritic):
             raise TypeError("Expecting 'actor_critic' to be an instance of ActorCritic")
 
-        # get policy
+        # get policy and value
         policy = actor_critic.actor
+        value = actor_critic.critic
 
         # create exploration strategy (wrap the original policy and specify how to explore)
         # By default, for discrete actions it will use a Categorical distribution and for continuous actions, it will
@@ -190,7 +191,10 @@ class PPO(GradientRLAlgo):
 
         # create estimator
         logger.debug('create return estimator (GAE)')
-        estimator = GAE(storage, gamma=gamma, tau=tau)
+        estimator = GAE(storage, value, gamma=gamma, tau=tau)
+
+        # create policy evaluator that will compute :math:`a \sim \pi(.|s_t)` and :math:`\pi(.|s_t)` on batch
+        policy_evaluator = PolicyEvaluator(policy=exploration)
 
         # create loss
         logger.debug('create loss')
@@ -204,7 +208,7 @@ class PPO(GradientRLAlgo):
         logger.debug('create explorer, evaluator, and updater')
         explorer = Explorer(task, exploration, storage, num_workers=num_workers)
         evaluator = Evaluator(estimator)
-        updater = Updater(policy, sampler, loss, optimizer)
+        updater = Updater(policy, sampler, loss, optimizer,  evaluators=[policy_evaluator])
 
         # initialize RL algorithm
         super(PPO, self).__init__(explorer, evaluator, updater)

pyrobolearn/algos/reinforce.py

@@ -149,7 +149,7 @@ class REINFORCE(GradientRLAlgo):
     - https://github.com/rlcode/reinforcement-learning/blob/master/2-cartpole/3-reinforce/cartpole_reinforce.py
     """
 
-    def __init__(self, task, approximators, gamma=0.99, lr=3e-4, num_workers=1):
+    def __init__(self, task, approximators, gamma=0.99, lr=0.001, num_workers=1):
         """
         Initialize the REINFORCE on-policy RL algorithm.
 
@@ -183,7 +183,7 @@ class REINFORCE(GradientRLAlgo):
             raise TypeError("Expecting the approximators to be an instance of `Policy`, or `ActorCritic`, instead got:"
                             " {}".format(type(approximators)))
 
-        # create exploration strategy
+        # create exploration strategy (if action is discrete, boltzmann exploration. If action is continuous, gaussian)
         exploration = ActionExploration(policy)
 
         # create storage
@@ -195,6 +195,7 @@ class REINFORCE(GradientRLAlgo):
         # create return: R_t = \sum_{t'=t}^{T} \gamma^{t'-t} r_{t'}
         returns = ActionRewardEstimator(storage, gamma=gamma)
 
+        # create policy evaluator that will compute :math:`a \sim \pi(.|s_t)` and :math:`\pi(.|s_t)` on batch
         policy_evaluator = PolicyEvaluator(policy=exploration)
 
         # create loss for policy: \mathbb{E}[ \log \pi_{\theta}(a_t | s_t) R_t ]

pyrobolearn/algos/rl_algo.py

@@ -10,6 +10,8 @@ Dependencies:
 import numpy as np
 # from pathos.multiprocessing import Pool
 
+from pyrobolearn.storages import RolloutStorage
+
 from pyrobolearn.algos.explorer import Explorer
 from pyrobolearn.algos.evaluator import Evaluator
 from pyrobolearn.algos.updater import Updater
@@ -246,6 +248,14 @@ class RLAlgo(object):  # Algo):
         """Return the storage unit."""
         return self.explorer.storage
 
+    @storage.setter
+    def storage(self, storage):
+        """Set the storage unit to the exploration, evaluation, and update phase."""
+        self.explorer.storage = storage
+        if self.evaluator is not None:
+            self.evaluator.storage = storage
+        self.updater.storage = storage
+
     @property
     def optimizers(self):
         """Return the optimizers."""
@@ -260,19 +270,28 @@ class RLAlgo(object):  # Algo):
     # Methods #
     ###########
 
-    def init(self, num_steps, num_rollouts, num_episodes, seed=None, *args, **kwargs):
+    def init(self, num_steps, num_rollouts, num_episodes, seed=None, verbose=False, *args, **kwargs):
         """
-        Initialize the reinforcement learning algorithm.
+        Initialize the reinforcement learning algorithm by creating and setting a new storage unit if necessary.
 
         Args:
             num_steps (int): number of step per rollout/trajectory
             num_rollouts (int): number of rollouts/trajectories per episode (default: 1)
             num_episodes (int): number of episodes (default: 1)
             seed (int): random seed
+            verbose (bool): if True, print details about the optimization process
             *args (list): list of optional arguments.
             **kwargs (dict): dictionary of optional arguments.
         """
-        pass
+        if isinstance(self.storage, RolloutStorage):
+            if self.storage.size != num_steps * num_rollouts:
+                if verbose:
+                    print("Creating new storage unit to have a new size with num_steps={} and num_rollouts={}"
+                          "".format(num_steps, num_rollouts))
+
+                state_shapes, action_shapes = self.storage.state_shapes, self.storage.action_shapes
+                self.storage.init(num_steps=num_steps, state_shapes=state_shapes, action_shapes=action_shapes,
+                                  num_trajectories=num_rollouts)
 
     # def init(self, explorer, evaluator, updater):
     #     """Initialize the RL algo."""
@@ -288,7 +307,7 @@ class RLAlgo(object):  # Algo):
             num_steps (int): number of step per rollout/trajectory
             num_rollouts (int): number of rollouts/trajectories per episode (default: 1)
             num_episodes (int): number of episodes (default: 1)
-            verbose(bool): if True, print details about the optimization process
+            verbose (bool): if True, print details about the optimization process
             seed (int): random seed
 
         Returns:
@@ -300,7 +319,7 @@ class RLAlgo(object):  # Algo):
             print("\n#### Start the RL algo ####")
 
         # init algo with the given parameters
-        self.init(num_steps=num_steps, num_rollouts=num_rollouts, num_episodes=num_episodes, seed=seed)
+        self.init(num_steps=num_steps, num_rollouts=num_rollouts, num_episodes=num_episodes, seed=seed, verbose=verbose)
 
         # set the policy in training mode
         self.policy.train()
@@ -321,7 +340,7 @@ class RLAlgo(object):  # Algo):
             # evaluate and update
             if self.evaluator is not None:
                 self.evaluator.evaluate(verbose=verbose)
-            losses = self.updater()
+            losses = self.updater.update(verbose=verbose)
 
             # add the loss in the history
             history.setdefault('losses', []).append(losses)
@@ -348,7 +367,7 @@ class RLAlgo(object):  # Algo):
         Returns:
             list: list of results for each policy at each time step
         """
-        results = self.task.run(self, num_steps=num_steps, dt=dt, use_terminating_condition=use_terminating_condition,
+        results = self.task.run(num_steps=num_steps, dt=dt, use_terminating_condition=use_terminating_condition,
                                 render=render)
         return results
 

pyrobolearn/algos/updater.py

@@ -124,10 +124,6 @@ class Updater(object):
     def storage(self, storage):
         """Set the storage unit."""
         self.sampler.storage = storage
-        # if not isinstance(storage, RolloutStorage):
-        #     raise TypeError("Expecting the storage to be an instance of `Storage`, instead got: "
-        #                     "{}".format(type(storage)))
-        # self._storage = storage
 
     @property
     def losses(self):

pyrobolearn/metrics/__init__.py

@@ -0,0 +1,12 @@
+
+# import the metrics
+from .metric import *
+
+# import reinforcement learning metrics
+from .rl_metrics import *
+
+# import imitation learning metrics
+from .il_metrics import *
+
+# import transfer learning metrics
+from .tl_metrics import *

pyrobolearn/metrics/il_metrics.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python
+"""Defines the metrics used in imitation learning (IL).
+"""
+
+import matplotlib.pyplot as plt
+
+from pyrobolearn.tasks import ILTask
+from pyrobolearn.metrics import Metric
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class ILMetric(Metric):
+    r"""Imitation Learning Metric
+
+    Metrics used in imitation learning.
+
+    References:
+        [1] "Learning from Humans", Billard et al., 2016
+    """
+
+    def __init__(self):
+        super(ILMetric, self).__init__()

pyrobolearn/metrics/metric.py

@@ -5,6 +5,10 @@ Dependencies:
 - `pyrobolearn.tasks`
 """
 
+import collections
+import matplotlib.pyplot as plt
+
+
 __author__ = "Brian Delhaisse"
 __copyright__ = "Copyright 2018, PyRoboLearn"
 __credits__ = ["Brian Delhaisse"]
@@ -24,45 +28,97 @@ class Metric(object):
     It notably contains the functionalities to evaluate a certain task using the metric, and different to plot them.
     """
 
-    def __init__(self):
+    def __init__(self, metrics=None):
+        """
+        Initialize the metric object.
+
+        Args:
+            metrics (None, Metric, list of Metric): inner metric objects.
+        """
+        self._metrics = metrics
+
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def metrics(self):
+        """Return the inner list of metric objects."""
+        return self._metrics
+
+    @metrics.setter
+    def metrics(self, metrics):
+        """Set the inner list of metrics."""
+        if metrics is None:
+            metrics = []
+        if not isinstance(metrics, collections.Iterable):
+            metrics = [metrics]
+        for i, metric in enumerate(metrics):
+            if not isinstance(metric, Metric):
+                raise TypeError("Expecting the given {}th metric to be an instance of `Metric`, instead got: "
+                                "{}".format(i, type(metric)))
+        self._metrics = metrics
+
+    ###########
+    # Methods #
+    ###########
+
+    def append(self, *args, **kwargs):
         pass
 
+    def update(self, *args, **kwargs):
+        pass
 
-class ILMetric(Metric):
-    r"""Imitation Learning Metric
+    def _plot(self, ax=None, filename=None):
+        """
+        Plot the metric. This has to be implemented in the child classes.
 
-    Metrics used in imitation learning.
+        Args:
+            ax (plt.Axes): axis to plot the figure.
+            filename (str, None): if a string is given, it will save the plot in the given filename.
 
-    References:
-        [1] "Learning from Humans", Billard et al., 2016
-    """
+        Returns:
+            plt.Axes: ax
+        """
+        pass
 
-    def __init__(self):
-        super(ILMetric, self).__init__()
+    def plot(self, ax=None, block=False, filename=None, subplots=()):
+        """
+        Plot the metric.
 
+        Args:
+            ax (plt.Axes): axis to plot the figure.
+            block (bool): if True, it will block when showing the graphs.
+            filename (str, None): if a string is given, it will save the plot in the given filename.
+        """
+        # if multiple metrics
+        if self.metrics:
 
-class RLMetric(Metric):
-    r"""Reinforcement Learning Metric
+            # if we want to use subplots
+            if len(subplots) > 0:
+                pass
 
-    Metrics used in reinforcement learning.
+            # if we just want multiple figures
+            for metric in self.metrics:
+                metric._plot(ax=ax, filename=filename)
 
-    References:
-        [1] "Deep Reinforcement Learning that Matters", Henderson et al., 2018
-    """
+            plt.show(block=block)
+        else:
+            self._plot(ax=ax, filename=filename)
+            plt.show(block=block)
 
-    def __init__(self):
-        super(RLMetric, self).__init__()
+    #############
+    # Operators #
+    #############
 
+    def __repr__(self):
+        """Return a representation string of the object."""
+        if self.metrics:
+            return ' + '.join(self.metrics)
+        return self.__class__.__name__
 
-class TLMetric(Metric):
-    r"""Transfer Learning Metric
-
-    Metrics used in transfer learning.
-
-    References:
-        [1] "A Survey on Transfer Learning", Pan et al., 2010
-        [2] "Transfer Learning for Reinforcement Learning Domains: A Survey", Taylor et al., 2009
-    """
-
-    def __init__(self):
-        super(TLMetric, self).__init__()
+    def __str__(self):
+        """Return a string describing the object."""
+        if self.metrics:
+            return ' + '.join(self.metrics)
+        return self.__class__.__name__

pyrobolearn/metrics/rl_metrics.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python
+"""Defines the metrics used in reinforcement learning (RL).
+"""
+
+import matplotlib.pyplot as plt
+
+from pyrobolearn.tasks import RLTask
+from pyrobolearn.algos import RLAlgo
+from pyrobolearn.metrics import Metric
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class RLMetric(Metric):
+    r"""Reinforcement Learning (abstract) Metric
+
+    Metrics used in reinforcement learning.
+
+    References:
+        [1] "Deep Reinforcement Learning that Matters", Henderson et al., 2018
+    """
+
+    def __init__(self):
+        """
+        Initialize the RL metric.
+        """
+        super(RLMetric, self).__init__()
+
+
+class AverageReturnMetric(RLMetric):
+    r"""Average / Expected return metric.
+
+    This computes the average / expected RL return given by:
+
+    .. math:: J(\pi_{\theta}) = \mathcal{E}_{\tau \sim \pi_{\theta}}[ R(\tau) ]
+
+    where :math:`R(\tau) = \sum_{t=0}^T \gamma^t r_t` is the discounted return.
+    """
+
+    def __init__(self, task, gamma=1.):
+        """
+        Initialize the average return metric.
+
+        Args:
+            gamma (float): discount factor
+        """
+        super(AverageReturnMetric, self).__init__()
+        self.gamma = gamma
+        self.task = task
+        self.returns = []
+
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def gamma(self):
+        """Return the discount factor"""
+        return self._gamma
+
+    @gamma.setter
+    def gamma(self, gamma):
+        """Set the discount factor"""
+        if gamma > 1.:
+            gamma = 1.
+        elif gamma < 0.:
+            gamma = 0.
+
+        self._gamma = gamma
+
+    @property
+    def task(self):
+        """Return the RL task."""
+        return self._task
+
+    @task.setter
+    def task(self, task):
+        """Set the RL task."""
+        if not isinstance(task, RLTask):
+            raise TypeError("Expecting the given 'task' to be an instance of `RLTask`, instead got: "
+                            "{}".format(type(task)))
+        self._task = task
+
+    ###########
+    # Methods #
+    ###########
+
+    def update(self):
+        pass
+
+    def _plot(self, ax=None, filename=None):
+        """
+        Plot the average return metric.
+
+        Args:
+            ax (plt.Axes): axis to plot the figure.
+            filename (str, None): if a string is given, it will save the plot in the given filename.
+        """
+        if ax is None:
+            fig, ax = plt.subplots()
+
+        ax.set_title('Average Return per iteration')  # per epoch, per iteration=epoch*batch
+        ax.set_xlabel('iterations')
+        ax.set_ylabel('Average return')
+        ax.plot(self.returns)
+
+        return ax
+
+
+class LossMetric(RLMetric):
+    r"""Loss Metric
+    """
+
+    def __init__(self, loss):
+        super(LossMetric, self).__init__()
+        self.loss = loss
+
+    def update(self):
+        pass

pyrobolearn/metrics/tl_metrics.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python
+"""Defines the metrics used in transfer learning (TL).
+"""
+
+import matplotlib.pyplot as plt
+
+from pyrobolearn.tasks import TLTask
+from pyrobolearn.metrics import Metric
+
+__author__ = "Brian Delhaisse"
+__copyright__ = "Copyright 2018, PyRoboLearn"
+__credits__ = ["Brian Delhaisse"]
+__license__ = "MIT"
+__version__ = "1.0.0"
+__maintainer__ = "Brian Delhaisse"
+__email__ = "briandelhaisse@gmail.com"
+__status__ = "Development"
+
+
+class TLMetric(Metric):
+    r"""Transfer Learning Metric
+
+    Metrics used in transfer learning.
+
+    References:
+        [1] "A Survey on Transfer Learning", Pan et al., 2010
+        [2] "Transfer Learning for Reinforcement Learning Domains: A Survey", Taylor et al., 2009
+    """
+
+    def __init__(self):
+        super(TLMetric, self).__init__()

pyrobolearn/policies/basic_policy.py

@@ -205,3 +205,17 @@ class PolicyFromQValue(Policy):
 
     # def sample(self, state):
     #     pass
+
+
+# Tests
+if __name__ == '__main__':
+    import copy
+    from pyrobolearn.states import FixedState
+    from pyrobolearn.actions import FixedAction
+
+    # check linear policy
+    policy = LinearPolicy(states=FixedState(range(4)), actions=FixedAction(range(2)))
+    print(policy)
+
+    target = copy.deepcopy(policy)
+    print(target)

pyrobolearn/states/state.py

@@ -2,7 +2,7 @@
 """Define the State class.
 
 This file defines the `State` class, which is returned by the environment, and given as an input to several
-models such as policies/controllers, dynamic transition functions, value estimators, reward/cost function, and so on.
+models such as policies/controllers, dynamic transition functions, value approximators, reward/cost function, and so on.
 """
 
 import numpy as np

pyrobolearn/storages/storage.py

@@ -639,6 +639,8 @@ class Batch(DictStorage):
         super(Batch, self).__init__(kwargs=kwargs, device=device, dtype=dtype, update=False)
         # contains the current values evaluated during the update phase of RL algorithms.
         self.current = DictStorage(kwargs={}, device=device, dtype=dtype, update=False)
+        # indices where the masks is different from 0 in the current batch
+        self.indices = None
 
     def get_current(self, key, default=None):
         """Try first to get the key from :attr:`current`, if not present, try to get it from the batch storage.
@@ -698,12 +700,12 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
             num_trajectories (int): number of trajectories.
         """
         # recurrent_hidden_state_size (int): size of the internal state
-        print("\nStorage: state shape: {}".format(state_shapes))
-        print("Storage: action shape: {}".format(action_shapes))
         super(RolloutStorage, self).__init__()
         self._step = np.zeros(int(num_trajectories), dtype=np.int)
         self._num_steps = int(num_steps)
         self._num_trajectories = int(num_trajectories)
+        self._state_shapes = state_shapes
+        self._action_shapes = action_shapes
         self._shifts = {}  # dictionary that maps the key to the time shift; this is add to the current time step
         self.init(self.num_steps, state_shapes, action_shapes, self.num_trajectories)
 
@@ -731,6 +733,16 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
         """Return the current time step."""
         return self._step
 
+    @property
+    def state_shapes(self):
+        """Return the state shapes."""
+        return self._state_shapes
+
+    @property
+    def action_shapes(self):
+        """Return the action shapes."""
+        return self._action_shapes
+
     ###########
     # Methods #
     ###########
@@ -822,6 +834,8 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
         self._step = np.zeros(int(num_trajectories), dtype=np.int)
         self._num_steps = int(num_steps)
         self._num_trajectories = int(num_trajectories)
+        self._state_shapes = state_shapes
+        self._action_shapes = action_shapes
 
         # allocate space for observations / states
         logger.debug('creating space for states with shape: {}'.format(state_shapes))
@@ -951,8 +965,6 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
             **kwargs (dict): dictionary containing other parameters to update in the storage. The other parameters
                 had to be added using the `create_new_entry()` method.
         """
-        print("Storage - insert state: {}".format(states))
-        print("Storage - insert action: {}".format(actions))
         t = self._step[rollout_idx]
 
         # check given observations/states and actions
@@ -1021,8 +1033,12 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
             else:  # value = tensor
                 batch[key] = sample(value, indices)
 
+        # create Batch object
+        batch = Batch(batch, device=self.device, dtype=self.dtype)
+        batch.indices = torch.tensor(range(len(indices)))[batch['masks'][:, 0] != 0].tolist()
+
         # return batch (which is given to the updater (and loss))
-        return Batch(batch, device=self.device, dtype=self.dtype)
+        return batch
 
     def end(self, rollout_idx=0, *args, **kwargs):
         """Once arrived at the end of an episode, it will fill the remaining mask values.

pyrobolearn/tasks/task.py

@@ -210,11 +210,7 @@ class Task(object):
         # results = []
         total_rewards = np.zeros(len(self.policies))
         self.reset()
-        # for t in range(4):
-        #     for policy in self.policies:
-        #         actions = policy.act(policy.states)
-        #     self.simulator.stepSimulation()
-        # time.sleep(2.)
+
         for t in count():
             if t >= num_steps:
                 break
@@ -264,7 +260,7 @@ class Task(object):
             return [policy.model for policy in self.policies]
         return self.policies[idx].model
 
-    def save_task(self, filename):
+    def save(self, filename):
         pass
 
 

pyrobolearn/values/value.py

@@ -265,8 +265,10 @@ class ParametrizedValue(ValueApproximator):
         # if the input is an instance of State, get the inner merged data.
         if isinstance(state, State):
             state = state.merged_data
-            if len(state) == 1:
-                state = state[0]
+
+        # if the input state is a list of len(1)
+        if isinstance(state, list) and len(state) == 1:
+            state = state[0]
 
         self.value = self.model.predict(state, to_numpy=to_numpy, return_logits=True, set_output_data=False)
         return self.value
@@ -437,13 +439,16 @@ class ParametrizedQValue(QValueApproximator):  # ParametrizedValue, QValueApprox
         # if the input is an instance of State, get the inner merged data.
         if isinstance(state, State):
             state = state.merged_data
-            if len(state) == 1:
-                state = state[0]
+        # if the input state is a list of len(1)
+        if isinstance(state, list) and len(state) == 1:
+            state = state[0]
 
+        # if the input is an insrtance of Action, get the inner merged data.
         if isinstance(action, Action):
             action = action.merged_data
-            if len(action) == 1:
-                action = action[0]
+        # if the input action is a list of len(1)
+        if isinstance(action, list) and len(action) == 1:
+            action = action[0]
 
         self.value = self.model.predict([state, action], to_numpy=to_numpy, return_logits=True, set_output_data=False)
         return self.value