<reponame>Karl-Krauth/robust-adaptive-lqr
"""adaptive.py

An object-oriented approach to implementing
different adaptive strategies.

"""


import numpy as np
import utils
import logging
import time
import itertools as it

from abc import ABC, abstractmethod


class AdaptiveMethod(ABC):
    """The base class for all adaptive methods

    The way to use this class is as follows:

    e = MyAdaptiveMethod(...)
    e.reset(rng)
    e.prime(num_iters, static_feedback, rng)
    for _ in range(horizon):
        cur_regret = e.step(rng)
        # do something with current regret
        # ...

    """

    def __init__(self, Q, R, A_star, B_star, sigma_w, rls_lam):
        """

        """

        # TODO: validate inputs

        self._Q = np.array(Q)
        self._R = np.array(R)
        self._A_star = np.array(A_star)
        self._B_star = np.array(B_star)
        self._sigma_w = sigma_w
        self._rls_lam = rls_lam
        self._n, self._p = B_star.shape

        if self._rls_lam is not None:
            self._rls = utils.RecursiveLeastSquaresEstimator(self._n, self._p, self._rls_lam)
        else:
            self._rls = None

        self._P_star, self._K_star = utils.dlqr(self._A_star, self._B_star, self._Q, self._R)
        self._J_star = (self._sigma_w ** 2) * np.trace(self._P_star)

    @abstractmethod
    def _design_controller(self, states, inputs, transitions, rng):
        """Design a controller for the next epoch.

        Called after an epoch is finished. Is handed two matrices
        of the history of states and inputs.

        returns a best estimate (Ahat, Bhat) of the true system.
        """

        pass

    @abstractmethod
    def _get_input(self, state, rng):
        """Obtain the next input to play from the current state

        """

        pass

    def _on_iteration_completion(self):
        """Called after an iteration is complete

        """

        pass

    def _on_epoch_completion(self):
        """Called after an epoch is complete

        """

        pass

    @abstractmethod
    def _should_terminate_epoch(self):
        """Return true if the epoch should terminate

        This method should not mutate state.

        """

        pass

    def reset(self, rng):
        """Reset both the dynamics and internal state.

        Must be called before ether prime() or step() is called.
        """
        self._state_history = []
        self._input_history = []
        self._transition_history = []
        self._cost_history = []

        # tracks the estimate errors for each epoch
        self._error_history = []

        # tracks the length of epochs
        self._epoch_history = []

        # tracks the average infinite time horizon cost
        self._infinite_horizon_cost_history = []

        if self._rls is not None:
            logger = self._get_logger()
            logger.debug("Using RLS estimator with rls_lam={}".format(self._rls_lam))
            self._rls = utils.RecursiveLeastSquaresEstimator(self._n, self._p, self._rls_lam)
        self._rls_history = []

        self._regret = 0
        self._epoch_idx = 0
        self._iteration_idx = 0
        self._iteration_within_epoch_idx = 0
        self._state_cur = np.zeros((self._n,))
        self._last_reset_time = time.time()
        self._has_primed = False

    def prime(self, num_iterations, static_feedback, excitation, rng):
        """Initialize the adaptive method with rollouts

        Must be called after reset() and before step() is called

        """

        assert num_iterations >= 1
        assert excitation > 0

        rng = self._get_rng(rng)

        for _ in range(num_iterations):
            inp = static_feedback.dot(self._state_cur) + excitation * rng.normal(size=(self._p,))
            xnext = self._A_star.dot(self._state_cur) + \
                    self._B_star.dot(inp) + \
                    self._sigma_w * rng.normal(size=(self._n,))

            self._state_history.append(self._state_cur)
            self._input_history.append(inp)
            self._transition_history.append(xnext)

            if self._rls is not None:
                self._rls.update(self._state_cur, inp, xnext)
                Ahat, Bhat, _ = self._rls.get_estimate()
                eps_A = np.linalg.norm(Ahat - self._A_star, ord=2)
                eps_B = np.linalg.norm(Bhat - self._B_star, ord=2)
                self._rls_history.append((eps_A, eps_B))

            self._state_cur = xnext

        Ahat, Bhat, Jnom = self._design_controller(
                np.array(self._state_history),
                np.array(self._input_history),
                np.array(self._transition_history),
                rng)
        eps_A = np.linalg.norm(Ahat - self._A_star, ord=2)
        eps_B = np.linalg.norm(Bhat - self._B_star, ord=2)

        logger = self._get_logger()
        logger.info("prime: eps_A={}, eps_B={}, Jnom={}".format(eps_A, eps_B, Jnom))

        self._error_history.append((eps_A, eps_B))
        self._infinite_horizon_cost_history.append(Jnom)

        # reset the initial state to zero
        self._state_cur = np.zeros_like(self._state_cur)

        self._has_primed = True


    def _get_rng(self, rng):
        return np.random if rng is None else rng

    def _get_logger(self):
        # sub-classes are encouraged to override this to provide a logger with
        # better name context
        return logging.getLogger(__name__)

    def step(self, rng):
        """Run the simulation forward.

        """

        try:
            self._state_history
        except AttributeError:
            raise ValueError("Call reset() before calling step()")

        if not self._has_primed:
            raise ValueError("Call prime() before calling step()")

        # get next input
        inp = self._get_input(self._state_cur, rng)

        # obtain current cost
        cost = utils.quad_form(self._Q, self._state_cur) + \
               utils.quad_form(self._R, inp)

        # advance to next state
        xnext = self._A_star.dot(self._state_cur) + \
                self._B_star.dot(inp) + \
                self._sigma_w * self._get_rng(rng).normal(size=(self._n,))

        self._state_history.append(self._state_cur)
        self._input_history.append(inp)
        self._transition_history.append(xnext)

        if self._rls is not None:
            self._rls.update(self._state_cur, inp, xnext)
            Ahat, Bhat, _ = self._rls.get_estimate()
            eps_A = np.linalg.norm(Ahat - self._A_star, ord=2)
            eps_B = np.linalg.norm(Bhat - self._B_star, ord=2)
            self._rls_history.append((eps_A, eps_B))

        self._cost_history.append(cost)
        self._state_cur = xnext

        self._regret += (cost - self._J_star)
        self._iteration_idx += 1
        self._iteration_within_epoch_idx += 1

        self._on_iteration_completion()

        if self._should_terminate_epoch():
            self.complete_epoch(rng)

        return self.regret()

    def complete_epoch(self, rng):
        """Call this method to forcefully complete an epoch

        """

        self._on_epoch_completion()

        # TODO(stephentu):
        # implement a recursive LS estimator of (A, B)
        # so we can report the error along the entire trajectory
        # instead of just at the epoch boundaries.
        Ahat, Bhat, infinite_horizon_cost = self._design_controller(
                np.array(self._state_history),
                np.array(self._input_history),
                np.array(self._transition_history),
                rng)
        eps_A = np.linalg.norm(Ahat - self._A_star, ord=2)
        eps_B = np.linalg.norm(Bhat - self._B_star, ord=2)
        self._error_history.append((eps_A, eps_B))
        self._infinite_horizon_cost_history.append(infinite_horizon_cost)

        logger = self._get_logger()
        logger.info("Finished with epoch {}, which lasted for {} out of {} iterations".format(
            self._epoch_idx,
            self._iteration_within_epoch_idx,
            self._iteration_idx))
        logger.info("Regret={}, eps_A={}, eps_B={}, Jhat={}, elapsed_time_since_reset={}".format(
            self.regret(),
            eps_A,
            eps_B,
            infinite_horizon_cost,
            time.time() - self._last_reset_time))

        self._epoch_idx += 1
        self._epoch_history.append(self._iteration_within_epoch_idx)
        self._iteration_within_epoch_idx = 0


    def regret(self):
        return self._regret

    def get_statistics(self, iteration_based):
        if iteration_based:
            def expand(h):
                assert len(h) == len(self._epoch_history) + 1
                return np.array(list(it.chain.from_iterable([[v]*repeat for v, repeat in zip(h[:-1], self._epoch_history)])))
            def rel_err(Jhat):
                return (Jhat - self._J_star)/self._J_star
            return (expand(self._error_history), expand([rel_err(e) for e in self._infinite_horizon_cost_history]))
        else:
            return (np.array(self._error_history), np.array(self._infinite_horizon_cost_history))

"""adaptive.py

An object-oriented approach to implementing
different adaptive strategies.

"""


import numpy as np
import utils
import logging
import time
import itertools as it

from abc import ABC, abstractmethod


class AdaptiveMethod(ABC):
    """The base class for all adaptive methods

    The way to use this class is as follows:

    e = MyAdaptiveMethod(...)
    e.reset(rng)
    e.prime(num_iters, static_feedback, rng)
    for _ in range(horizon):
        cur_regret = e.step(rng)
        # do something with current regret
        # ...

    """

    def __init__(self, Q, R, A_star, B_star, sigma_w, rls_lam):
        """

        """

        # TODO: validate inputs

        self._Q = np.array(Q)
        self._R = np.array(R)
        self._A_star = np.array(A_star)
        self._B_star = np.array(B_star)
        self._sigma_w = sigma_w
        self._rls_lam = rls_lam
        self._n, self._p = B_star.shape

        if self._rls_lam is not None:
            self._rls = utils.RecursiveLeastSquaresEstimator(self._n, self._p, self._rls_lam)
        else:
            self._rls = None

        self._P_star, self._K_star = utils.dlqr(self._A_star, self._B_star, self._Q, self._R)
        self._J_star = (self._sigma_w ** 2) * np.trace(self._P_star)

    @abstractmethod
    def _design_controller(self, states, inputs, transitions, rng):
        """Design a controller for the next epoch.

        Called after an epoch is finished. Is handed two matrices
        of the history of states and inputs.

        returns a best estimate (Ahat, Bhat) of the true system.
        """

        pass

    @abstractmethod
    def _get_input(self, state, rng):
        """Obtain the next input to play from the current state

        """

        pass

    def _on_iteration_completion(self):
        """Called after an iteration is complete

        """

        pass

    def _on_epoch_completion(self):
        """Called after an epoch is complete

        """

        pass

    @abstractmethod
    def _should_terminate_epoch(self):
        """Return true if the epoch should terminate

        This method should not mutate state.

        """

        pass

    def reset(self, rng):
        """Reset both the dynamics and internal state.

        Must be called before ether prime() or step() is called.
        """
        self._state_history = []
        self._input_history = []
        self._transition_history = []
        self._cost_history = []

        # tracks the estimate errors for each epoch
        self._error_history = []

        # tracks the length of epochs
        self._epoch_history = []

        # tracks the average infinite time horizon cost
        self._infinite_horizon_cost_history = []

        if self._rls is not None:
            logger = self._get_logger()
            logger.debug("Using RLS estimator with rls_lam={}".format(self._rls_lam))
            self._rls = utils.RecursiveLeastSquaresEstimator(self._n, self._p, self._rls_lam)
        self._rls_history = []

        self._regret = 0
        self._epoch_idx = 0
        self._iteration_idx = 0
        self._iteration_within_epoch_idx = 0
        self._state_cur = np.zeros((self._n,))
        self._last_reset_time = time.time()
        self._has_primed = False

    def prime(self, num_iterations, static_feedback, excitation, rng):
        """Initialize the adaptive method with rollouts

        Must be called after reset() and before step() is called

        """

        assert num_iterations >= 1
        assert excitation > 0

        rng = self._get_rng(rng)

        for _ in range(num_iterations):
            inp = static_feedback.dot(self._state_cur) + excitation * rng.normal(size=(self._p,))
            xnext = self._A_star.dot(self._state_cur) + \
                    self._B_star.dot(inp) + \
                    self._sigma_w * rng.normal(size=(self._n,))

            self._state_history.append(self._state_cur)
            self._input_history.append(inp)
            self._transition_history.append(xnext)

            if self._rls is not None:
                self._rls.update(self._state_cur, inp, xnext)
                Ahat, Bhat, _ = self._rls.get_estimate()
                eps_A = np.linalg.norm(Ahat - self._A_star, ord=2)
                eps_B = np.linalg.norm(Bhat - self._B_star, ord=2)
                self._rls_history.append((eps_A, eps_B))

            self._state_cur = xnext

        Ahat, Bhat, Jnom = self._design_controller(
                np.array(self._state_history),
                np.array(self._input_history),
                np.array(self._transition_history),
                rng)
        eps_A = np.linalg.norm(Ahat - self._A_star, ord=2)
        eps_B = np.linalg.norm(Bhat - self._B_star, ord=2)

        logger = self._get_logger()
        logger.info("prime: eps_A={}, eps_B={}, Jnom={}".format(eps_A, eps_B, Jnom))

        self._error_history.append((eps_A, eps_B))
        self._infinite_horizon_cost_history.append(Jnom)

        # reset the initial state to zero
        self._state_cur = np.zeros_like(self._state_cur)

        self._has_primed = True


    def _get_rng(self, rng):
        return np.random if rng is None else rng

    def _get_logger(self):
        # sub-classes are encouraged to override this to provide a logger with
        # better name context
        return logging.getLogger(__name__)

    def step(self, rng):
        """Run the simulation forward.

        """

        try:
            self._state_history
        except AttributeError:
            raise ValueError("Call reset() before calling step()")

        if not self._has_primed:
            raise ValueError("Call prime() before calling step()")

        # get next input
        inp = self._get_input(self._state_cur, rng)

        # obtain current cost
        cost = utils.quad_form(self._Q, self._state_cur) + \
               utils.quad_form(self._R, inp)

        # advance to next state
        xnext = self._A_star.dot(self._state_cur) + \
                self._B_star.dot(inp) + \
                self._sigma_w * self._get_rng(rng).normal(size=(self._n,))

        self._state_history.append(self._state_cur)
        self._input_history.append(inp)
        self._transition_history.append(xnext)

        if self._rls is not None:
            self._rls.update(self._state_cur, inp, xnext)
            Ahat, Bhat, _ = self._rls.get_estimate()
            eps_A = np.linalg.norm(Ahat - self._A_star, ord=2)
            eps_B = np.linalg.norm(Bhat - self._B_star, ord=2)
            self._rls_history.append((eps_A, eps_B))

        self._cost_history.append(cost)
        self._state_cur = xnext

        self._regret += (cost - self._J_star)
        self._iteration_idx += 1
        self._iteration_within_epoch_idx += 1

        self._on_iteration_completion()

        if self._should_terminate_epoch():
            self.complete_epoch(rng)

        return self.regret()

    def complete_epoch(self, rng):
        """Call this method to forcefully complete an epoch

        """

        self._on_epoch_completion()

        # TODO(stephentu):
        # implement a recursive LS estimator of (A, B)
        # so we can report the error along the entire trajectory
        # instead of just at the epoch boundaries.
        Ahat, Bhat, infinite_horizon_cost = self._design_controller(
                np.array(self._state_history),
                np.array(self._input_history),
                np.array(self._transition_history),
                rng)
        eps_A = np.linalg.norm(Ahat - self._A_star, ord=2)
        eps_B = np.linalg.norm(Bhat - self._B_star, ord=2)
        self._error_history.append((eps_A, eps_B))
        self._infinite_horizon_cost_history.append(infinite_horizon_cost)

        logger = self._get_logger()
        logger.info("Finished with epoch {}, which lasted for {} out of {} iterations".format(
            self._epoch_idx,
            self._iteration_within_epoch_idx,
            self._iteration_idx))
        logger.info("Regret={}, eps_A={}, eps_B={}, Jhat={}, elapsed_time_since_reset={}".format(
            self.regret(),
            eps_A,
            eps_B,
            infinite_horizon_cost,
            time.time() - self._last_reset_time))

        self._epoch_idx += 1
        self._epoch_history.append(self._iteration_within_epoch_idx)
        self._iteration_within_epoch_idx = 0


    def regret(self):
        return self._regret

    def get_statistics(self, iteration_based):
        if iteration_based:
            def expand(h):
                assert len(h) == len(self._epoch_history) + 1
                return np.array(list(it.chain.from_iterable([[v]*repeat for v, repeat in zip(h[:-1], self._epoch_history)])))
            def rel_err(Jhat):
                return (Jhat - self._J_star)/self._J_star
            return (expand(self._error_history), expand([rel_err(e) for e in self._infinite_horizon_cost_history]))
        else:
            return (np.array(self._error_history), np.array(self._infinite_horizon_cost_history))
adaptive.py  An object-oriented approach to implementing different adaptive strategies. The base class for all adaptive methods      The way to use this class is as follows:      e = MyAdaptiveMethod(...)     e.reset(rng)     e.prime(num_iters, static_feedback, rng)     for _ in range(horizon):         cur_regret = e.step(rng)         # do something with current regret         # ... # TODO: validate inputs Design a controller for the next epoch.          Called after an epoch is finished. Is handed two matrices         of the history of states and inputs.          returns a best estimate (Ahat, Bhat) of the true system. Obtain the next input to play from the current state Called after an iteration is complete Called after an epoch is complete Return true if the epoch should terminate          This method should not mutate state. Reset both the dynamics and internal state.          Must be called before ether prime() or step() is called. # tracks the estimate errors for each epoch # tracks the length of epochs # tracks the average infinite time horizon cost Initialize the adaptive method with rollouts          Must be called after reset() and before step() is called # reset the initial state to zero # sub-classes are encouraged to override this to provide a logger with # better name context Run the simulation forward. # get next input # obtain current cost # advance to next state Call this method to forcefully complete an epoch # TODO(stephentu): # implement a recursive LS estimator of (A, B) # so we can report the error along the entire trajectory # instead of just at the epoch boundaries.