models.py

"""
This project was developed by Rocky Duan, Peter Chen, Pieter Abbeel for the Berkeley Deep RL Bootcamp, August 2017. Bootcamp website with slides and lecture videos: https://sites.google.com/view/deep-rl-bootcamp/.

Copyright 2017 Deep RL Bootcamp Organizers.

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

from collections import OrderedDict

from utils import *
import numpy as np
import chainer
import chainer.links as L
import chainer.functions as F
from chainer import Variable
import scipy.optimize
import gym.spaces


class Model(chainer.Chain):
    def __init__(self, observation_space, action_space, env_spec, **kwargs):
        super().__init__(**kwargs)
        self.observation_space = observation_space
        self.action_space = action_space
        self.env_spec = env_spec

        self.obs_dim = flatten_dim(observation_space)
        self.action_dim = flatten_dim(action_space)


class NNFeatureModel(Model):
    feature_dim = None

    def compute_features(self, obs):
        raise NotImplementedError

    def feature_links(self):
        raise NotImplementedError


class MLPFeatureModel(NNFeatureModel):
    def __init__(self, observation_space, action_space, env_spec, *, hidden_sizes=(128, 64), hidden_nonlinearity=F.relu,
                 **kwargs):
        super().__init__(observation_space, action_space, env_spec, **kwargs)
        self.hidden_sizes = hidden_sizes
        if isinstance(hidden_nonlinearity, str):
            if hidden_nonlinearity == 'relu':
                hidden_nonlinearity = F.relu
            elif hidden_nonlinearity == 'tanh':
                hidden_nonlinearity = F.tanh
            elif hidden_nonlinearity == 'elu':
                hidden_nonlinearity = F.elu
            else:
                raise NotImplementedError
        self.hidden_nonlinearity = hidden_nonlinearity
        self.n_layers = len(hidden_sizes)
        self._feature_links = OrderedDict()
        with self.init_scope():
            input_size = self.obs_dim
            for idx, hidden_size in enumerate(hidden_sizes):
                link = L.Linear(input_size, hidden_size)
                name = "fc{}".format(idx + 1)
                setattr(self, name, link)
                self._feature_links[name] = link
                input_size = hidden_size
            self.feature_dim = input_size

    def feature_links(self):
        return self._feature_links

    def compute_features(self, obs):
        obs = F.cast(obs, np.float32)
        h = obs
        for link in self.feature_links().values():
            h = self.hidden_nonlinearity(link(h))
        return h


class CNNFeatureModel(NNFeatureModel):
    def __init__(self, observation_space, action_space, env_spec, **kwargs):
        super().__init__(observation_space, action_space, env_spec, **kwargs)

        with self.init_scope():
            in_channels = observation_space.shape[-1]
            self.conv1 = L.Convolution2D(in_channels, 16, 8, stride=4)
            self.conv2 = L.Convolution2D(16, 32, 4, stride=2)
            self.fc = L.Linear(2592, 256)
            self.feature_dim = 256

    def feature_links(self):
        return dict(conv1=self.conv1, conv2=self.conv2, fc=self.fc)

    def compute_features(self, obs):
        obs = F.cast(obs, np.float32)
        obs = F.transpose(obs, (0, 3, 1, 2))
        h1 = F.relu(self.conv1(obs))
        h2 = F.relu(self.conv2(h1))
        h3 = F.relu(self.fc(h2))
        return h3


class WeightSharingFeatureModel(NNFeatureModel):
    def __init__(self, observation_space, action_space, env_spec, feature_model, **kwargs):
        super().__init__(observation_space, action_space, env_spec, **kwargs)
        with self.init_scope():
            for name, link in feature_model.feature_links().items():
                setattr(self, name, link)
        self._compute_features = feature_model.compute_features
        self.feature_dim = feature_model.feature_dim
        self.feature_model = feature_model
        self._feature_links = feature_model.feature_links

    def compute_features(self, obs):
        return self._compute_features(obs)

    def feature_links(self):
        return self._feature_links()


# ==============================
# Value Functions
# ==============================

class ValueFunction(Model):
    def compute_state_values(self, obs):
        raise NotImplementedError


class NNFeatureValueFunction(ValueFunction, NNFeatureModel):
    def __init__(self, observation_space, action_space, env_spec, **kwargs):
        super().__init__(observation_space, action_space, env_spec, **kwargs)
        with self.init_scope():
            self.l_vf = L.Linear(self.feature_dim, 1)

    def compute_state_values(self, obs, feats=None):
        if feats is None:
            feats = super().compute_features(obs)
        return self.l_vf(feats)[..., 0]


class ZeroValueFunction(ValueFunction):
    def compute_state_values(self, obs):
        obs_data = obs.data
        xp = chainer.cuda.get_array_module(obs_data)
        return chainer.Variable(xp.zeros(len(obs_data), dtype=xp.float32))

    pass


class WeightSharingValueFunction(NNFeatureValueFunction, WeightSharingFeatureModel):
    pass


class MLPValueFunction(NNFeatureValueFunction, MLPFeatureModel):
    pass


class CNNValueFunction(NNFeatureValueFunction, CNNFeatureModel):
    pass


# ==============================
# Policies
# ==============================

class Policy(Model):
    feature_dim = None

    def compute_dists(self, obs):
        """
        Given some observations, compute the parameters for the action distributions.
        :param obs: A chainer variable containing a list of observations.
        :return: An instance of the Distribution class, represeting the action distribution.
        """
        raise NotImplementedError

    def get_actions(self, obs):
        with chainer.no_backprop_mode():
            dists = self.compute_dists(Variable(np.asarray(obs)))
            actions = dists.sample()
            return actions.data, {k: v.data for k, v in dists.as_dict().items()}

    def get_action(self, ob):
        actions, dists = self.get_actions(np.expand_dims(ob, 0))
        return actions[0], {k: v[0] for k, v in dists.items()}


class NormalizingPolicy(Policy):
    """
    This policy wraps an existing policy. It mains a running average of past observations,
    to roughly ensure that the input dimensions are
    whitened.
    Most useful for relatively low-dimensional observations (i.e. not images)
    """

    def __init__(self, policy, clip=None):
        super().__init__(policy.observation_space, policy.action_space, policy.env_spec)
        with self.init_scope():
            self.policy = policy
        self.running_stat = RunningStat(policy.observation_space.shape)
        self.clip = float(clip)

    def get_actions(self, obs):
        obs = np.asarray(obs)
        for ob in obs:
            self.running_stat.push(ob)
        obs = (obs - self.running_stat.mean[None, ...]) / \
            (self.running_stat.std[None, ...] + 1e-8)
        if self.clip is not None:
            obs = np.clip(obs, -self.clip, self.clip)
        return self.policy.get_actions(obs)

    def compute_dists(self, obs):
        mean_var = Variable(self.running_stat.mean.astype(np.float32))
        std_var = Variable(self.running_stat.std.astype(np.float32))

        obs = obs - F.broadcast_to(mean_var, obs.shape)
        obs = obs / (F.broadcast_to(std_var, obs.shape) + 1e-8)

        if self.clip is not None:
            obs = F.clip(obs, -self.clip, self.clip)

        return self.policy.compute_dists(obs)

    @property
    def distribution(self):
        return self.policy.distribution


class NNFeaturePolicy(Policy, NNFeatureModel):
    def create_vf(self):
        return WeightSharingValueFunction(
            observation_space=self.observation_space,
            action_space=self.action_space,
            env_spec=self.env_spec,
            feature_model=self,
        )


class CategoricalPolicy(NNFeaturePolicy):
    def __init__(self, observation_space, action_space, env_spec, **kwargs):
        super().__init__(observation_space, action_space, env_spec, **kwargs)
        with self.init_scope():
            assert self.feature_dim is not None
            self.l_act = L.Linear(self.feature_dim, self.action_dim)

    def compute_dists(self, obs, feats=None):
        if feats is None:
            feats = super().compute_features(obs)
        logits = self.l_act(feats)
        return Categorical(logits=logits)

    @property
    def distribution(self):
        return Categorical


class GaussianPolicy(NNFeaturePolicy):
    def __init__(self, observation_space, action_space, env_spec, **kwargs):
        super().__init__(observation_space, action_space, env_spec, **kwargs)
        with self.init_scope():
            assert self.feature_dim is not None
            self.l_act = L.Linear(self.feature_dim, self.action_dim)
            self.log_std = chainer.Parameter(
                shape=(self.action_dim,), initializer=chainer.initializers.Zero())

    def compute_dists(self, obs, feats=None):
        if feats is None:
            feats = super().compute_features(obs)
        means = self.l_act(feats)
        # for this policy, the variance is independent of the state
        log_stds = F.tile(self.log_std.reshape((1, -1)), (len(feats), 1))
        return Gaussian(means=means, log_stds=log_stds)

    @property
    def distribution(self):
        return Gaussian


class CategoricalMLPPolicy(CategoricalPolicy, MLPFeatureModel):
    pass


class GaussianMLPPolicy(GaussianPolicy, MLPFeatureModel):
    pass


class CategoricalCNNPolicy(CategoricalPolicy, CNNFeatureModel):
    pass


# ==============================
# Baselines
# ==============================


class Baseline(object):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def predict(self, obs):
        raise NotImplementedError

    def update(self, trajs):
        raise NotImplementedError


class ZeroBaseline(Baseline):
    def predict(self, obs):
        return np.zeros(len(obs))

    def update(self, trajs):
        pass


class TimeDependentBaseline(Baseline):
    def __init__(self, observation_space, action_space, env_spec):
        self.observation_space = observation_space
        self.action_space = action_space
        self.env_spec = env_spec
        # 1D array: b[t] = moving average of past returns R_t
        self.b = np.zeros(env_spec.timestep_limit + 1)
        # # 1D array: ns[t] = number of past trajectories with at least t time steps
        # self.ns = np.zeros(env_spec.timestep_limit)

    def predict(self, obs):
        T = obs.shape[0]
        return self.b[:T]

    def update(self, trajs):
        all_returns = [traj['returns'] for traj in trajs]
        max_T = self.env_spec.timestep_limit + 1
        padded_returns = [
            np.pad(returns, (0, max_T - len(returns)),
                   mode='constant', constant_values=(np.nan,))
            for returns in all_returns
        ]
        self.b[:] = np.nanmean(padded_returns, axis=0)
        self.b[np.isnan(self.b)] = 0


class LinearFeatureBaseline(Baseline):
    def __init__(self, observation_space, action_space, env_spec):
        self.observation_space = observation_space
        self.action_space = action_space
        self.env_spec = env_spec
        self.coeffs = None

    def get_features(self, obs):
        o = np.clip(obs, -10, 10)
        l = len(obs)
        al = np.arange(l).reshape(-1, 1) / 100.0
        return np.concatenate([o, o ** 2, al, al ** 2, al ** 3, np.ones((l, 1))], axis=1)

    def update(self, trajs):
        featmat = np.concatenate(
            [self.get_features(traj["observations"]) for traj in trajs])
        returns = np.concatenate([traj["returns"] for traj in trajs])
        self.coeffs = np.linalg.lstsq(
            featmat.T.dot(featmat), featmat.T.dot(returns))[0]

    def predict(self, obs):
        if self.coeffs is None:
            return np.zeros(len(obs))
        return self.get_features(obs).dot(self.coeffs)


class NNFeatureBaseline(Baseline, NNFeatureModel):
    def __init__(self, observation_space, action_space, env_spec, mixture_fraction=0.1, **kwargs):
        # For input, we will concatenate the observation with the time, so we need to increment the observation
        # dimension
        if isinstance(observation_space, gym.spaces.Box) and len(observation_space.shape) == 1:
            self.concat_time = True
            obs_space = gym.spaces.Box(
                low=np.append(observation_space.low, 0),
                high=np.append(observation_space.high, 2 ** 32),
            )
        else:
            obs_space = observation_space
        self.mixture_fraction = mixture_fraction
        super().__init__(obs_space, action_space, env_spec, **kwargs)
        with self.init_scope():
            assert self.feature_dim is not None
            self.l_b = L.Linear(self.feature_dim, self.action_dim)

    def compute_baselines(self, obs):
        feats = self.compute_features(obs)
        return self.l_b(feats)[..., 0]

    # Value functions can themselves be used as baselines
    def predict(self, obs):
        with chainer.no_backprop_mode():
            if self.concat_time:
                ts = np.arange(len(obs)) / self.env_spec.timestep_limit
                obs = np.concatenate([obs, ts[:, None]], axis=-1)
            values = self.compute_baselines(Variable(obs))
            return values.data

    # By default, when used as baselines, value functions are updated via L-BFGS
    def update(self, trajs):
        obs = np.concatenate([traj['observations'] for traj in trajs], axis=0)
        if self.concat_time:
            ts = np.concatenate([np.arange(len(traj['observations'])) / self.env_spec.timestep_limit for traj in trajs],
                                axis=0)
            obs = np.concatenate([obs, ts[:, None]], axis=-1)
        returns = np.concatenate([traj['returns'] for traj in trajs], axis=0)
        baselines = np.concatenate([traj['baselines']
                                    for traj in trajs], axis=0)

        # regress to a mixture of current and past predictions
        targets = returns * (1. - self.mixture_fraction) + \
            baselines * self.mixture_fraction

        # use lbfgs to perform the update
        cur_params = get_flat_params(self)

        obs = Variable(obs)
        targets = Variable(targets.astype(np.float32))

        def f_loss_grad(x):
            set_flat_params(self, x)
            self.cleargrads()
            values = self.compute_baselines(obs)
            loss = F.mean(F.square(values - targets))
            loss.backward()
            flat_grad = get_flat_grad(self)
            return loss.data.astype(np.float64), flat_grad.astype(np.float64)

        new_params = scipy.optimize.fmin_l_bfgs_b(
            f_loss_grad, cur_params, maxiter=10)[0]

        set_flat_params(self, new_params)


class MLPBaseline(NNFeatureBaseline, MLPFeatureModel):
    pass