%!TEX root = ../notes.tex
\section{February 3, 2022}

\subsection{Arithmetic Functions}

We look at arithmetic functions and how they act on prime numbers:
\begin{definition}[Arithmetic Function]
    An \ul{arithmetic function} is a function $f: \ZZ_+ \to \CC$.

    \emph{(Typically, these are integer valued.)}
\end{definition}
\begin{example}We have some examples of arithmetic functions:
    \begin{itemize}
        \item Euler $\phi$ function.
        \item $\tau(n)$, the counting function. It takes a positive integer and counts the number of positive divisors of $n$. \[\tau(n) = \sum_{d\mid n} 1\]
        \item $\sigma(n)$, the sum of divisors function. It is the sum over all the positive divisors of $n$. \[\sigma(n) = \sum_{d\mid n}d\]
    \end{itemize}
\end{example}
We have some properties of these functions, like \emph{multiplicative}, \emph{completely multiplicative}, \emph{additive}, \emph{completely additive}.

\begin{definition}[Multiplicativity]
    An arithmetic function $f$ is \ul{multiplicative} if
    \begin{align*}
        f(mn) & = f(m)f(n)\quad \text{whenever }(m, n) = 1
        \intertext{$f$ is said to be \ul{totally or completely multiplicative} if }
        f(mn) & = f(m)f(n)\quad \forall m, n\in\ZZ_+
    \end{align*}
    regardless of coprimality.
\end{definition}

If $f$ is multiplicative and $n_1, \dots, n_k$ are positive pairwise coprime integers, then \[f(n_1\dots n_k) = f(n_1)f(n_2)\dots f(n_k).\]

A particular case that is useful is when we write
\[n=p_1^{e_1}p_2^{e_2}\cdots p_k^{e_k}\]
so that assuming multiplicativity, we have that
\[f(n) = f(p_1^{e_1}) f(p_2^{e_2}) \cdots f(p_k^{e_k})\]
A common type of arithmetic function is a summatory function, namely a function $f$ of the form
\[f(n) = \sum_{d\mid n} g(d), \quad \text{where $g$ is some arithmetic function.}\]

\emph{Food for thought:} how special are summatory functions within the set of all arithmetic functions?

A special property of summatory functions is that they ``inherit multiplicativity''.
\begin{lemma}\label{lemma:summatory-preserves-multiplicativity}
    If $g$ is a multiplicative function, and \[f(n) = \sum_{d\mid n}g(d)\quad \forall n,\] then $f$ itself is multiplicative.
\end{lemma}

\begin{proof}
    Suppose $m, n\in \ZZ_+$ are coprime positive integers.

    The divisors $d$ of $mn$ are the products $a\cdot b$ where $a\mid m$ and $b\mid n$. Each such pair $a, b$ yields a uniquely determined produce $d = a\cdot b$. Conversely, since $(m, n) = 1$, each divisor $d$ of $mn$ determines a unique divisor $a=\gcd(d, m)$ and $b = \gcd(d, n)$ so that $d = a\cdot b$.

    Thus there is a \emph{bijection} between divisors of $mn$ and $m, n$ separately
    \[d\mid mn \longleftrightarrow (a\mid m, b\mid n)\]
    Thus we have
    \begin{align*}
        f(m\cdot n) & = \sum_{d\mid mn}g(d)                                              \\
                    & = \sum_{a\mid m}\sum_{b\mid n}g(ab)                                \\
                    & = \sum_{a\mid m}\sum_{b\mid n}g(a)g(b)                             \\
                    & = \left(\sum_{a\mid m} g(a)\right) \left(\sum_{b\mid n}g(b)\right)
        = f(m)\cdot f(n)
    \end{align*}
    Thus completes the proof that $f$ is multiplicative.
\end{proof}

\recall The functions introduced earlier
\[\tau(n) = \sum_{d\mid n} 1\qquad \sigma(n) = \sum_{d\mid n}d\]
So $\tau$ is the summatory function of the constant $1$ functions, and $\sigma$ is the summatory function of the identity function. We know that the constant $1$ function and the identity function are both completely multiplicative, so $\sigma$ and $\tau$ are multiplicative functions.

The implication of which is that it suffices to apply $\tau$ and $\sigma$ on prime powers and multiply.

Let $p$ be a prime. Then
\begin{align*}
    \tau(p^e)   & = e + 1\quad  \text{(from $p^0$ to $p^e$)}.          \\
    \intertext{We also have}
    \sigma(p^e) & = 1 + p + p^2 + \cdots + p^e = \frac{p^{e+1}-1}{p-1}
\end{align*}
Therefore, if $n=p_1^{e_1}p_2^{e_2}\cdots p_k^{e_k}$, then
\begin{align*}
    \tau(n)   & = \prod_{i=1}^k (e_i + 1)                                 \\
    \sigma(n) & = \prod_{i=1}^k \left(\frac{p_i^{e_i+1}-1}{p_i-1}\right).
\end{align*}
\begin{remark}
    There are higher order divisor functions
    \[\sigma_k(n) = \sum_{d\mid n}d^k\]
    so $\sigma_0 = \tau, \sigma_1 = \sigma, \dots$
\end{remark}

\subsection{Review of \texorpdfstring{$\ZZ/n\ZZ$}{Z/nZ} and its units}

\begin{definition}[Modular Congruence]
    If $a, b, m\in \ZZ$, $m\neq 0$, we say that \ul{$a$ is congruent to $b$ modulo $m$} if $m\mid b-a$. We write
    \begin{equation*}
        a\equiv b\mod{m}, \text{ or more simply } a\equiv b\ (m)
    \end{equation*}
\end{definition}
Congruence mod $m$ is an equivalence relation on $\ZZ$. If $a\in \ZZ$, $\overline{a}$ denotes the set of integers congruent to $a\mod m$, i.e. $\overline{a} = \{a + km \mid k\in \ZZ\}$.

\begin{definition}[$\ZZ/m\ZZ$, Residues mod $m$]
    The set of congruence classes mod $m$ is denoted $\ZZ/m\ZZ$. This is a quotient ring of the ring of integers $\ZZ$.

    If $\overline{a}_1, \overline{a}_2, \dots, \overline{a}_m$ form a complete set of congruence classes mod $m$, then the set of integers $\{a_1, a_2, \dots, a_m\}$ is called a \ul{complete set of residues mod $m$}.
\end{definition}

$\ZZ/m\ZZ$ can be endowed with the structure of a commutative ring by setting
\begin{align*}
    \overline{a} + \overline{b}               & = \overline{a + b} \\
    \text{and }\overline{a}\cdot \overline{b} & = \overline{ab},
\end{align*}
and proving that this is well-defined as ring operations.

\begin{proposition}
    The set of units in $\ZZ/m\ZZ$ is exactly
    \[\{\overline{a}\mid (a, m) = 1\}\]
\end{proposition}
\begin{proof}
    Let $\overline{a}\in \ZZ/m\ZZ$, then
    \begin{align*}
                              & \exists \overline{b}\in \ZZ/m\ZZ\text{ s.t. }\overline{b}\cdot \overline{a}\equiv 1\mod m \\
        \Longleftrightarrow\  & \exists b, n\in \ZZ\text{ s.t. } ba - mn = 1                                              \\
        \intertext{Then by B\'ezout's identity...}
        \Longleftrightarrow\  & (a, m) = 1
    \end{align*}
\end{proof}

\subsection{The Euler \texorpdfstring{$\phi$}{phi} Function}
For $n\in \ZZ_+$, $\phi(n)$ is defined to be the number of integers $1\leq m\leq n$ coprime to $n$.
\begin{example}
    We have some examples of the Euler $\phi$ functions:
    \begin{align*}
        \phi(1)   & = 1                                                                 \\
        \phi(p)   & = p-1 \text{ for any prime $p$}
        \intertext{Let $e\geq 1$, }
        \phi(p^e) & = p^e-p^{e-1} \text{ for prime powers, we exclude multiples of $p$}
    \end{align*}
\end{example}

\emph{Wouldn't be great if $\phi$ were multiplicative? It is!}
\begin{theorem}
    If $(m,n) = 1$, then $\phi(mn) = \phi(m)\phi(n)$.
\end{theorem}
\begin{proof}
    By the Chinese Remainder Theorem\footnote{This is an easy way to prove this assuming Math 1530 (Abstract Algebra). There is another way to prove this with one hand tied behind the back, it just takes more mental muscle to do. }, $\ZZ/mn\ZZ\cong \ZZ/m\ZZ\times \ZZ/n\ZZ$ if $(m, n) = 1$.

    Taking the unit groups on both sides, we have
    \[(\ZZ/mn\ZZ)^\times \cong (\ZZ/m\ZZ)^\times \times (\ZZ/n\ZZ)^\times\]
    and the Euler $\phi$ function is simply measuring the order of said unit groups ($\phi(n) = \left|(\ZZ/n\ZZ)^\times\right|$).
\end{proof}

Here is an important fact about the Euler $\phi$ function:
\begin{proposition}
    We have
    \begin{equation*}
        \sum_{d\mid n}\phi(d) = n.
    \end{equation*}
\end{proposition}
\begin{proof}\emph{(1: a cute, snazzy proof)}
    Consider the $n$ rational numbers
    \[\frac{1}{n}, \frac{2}{n}, \dots, \frac{n-1}{n}, \frac{n}{n} = 1\]
    and reduce all to lowest terms so that the numerator and denominator are coprime.

    Q: Given a positive divisor $d$ of $n$, how many fractions have $d$ as the denominator?

    A: We have exactly $\phi(d)$ of them.

    Conversely, every denominator $d$ is certainly a divisor of $n$. So we conclude that $\displaystyle n = \sum_{d\mid n}\phi(d)$.
\end{proof}
\begin{proof}\emph{(2: using what we've learnt)}
    We use the fact that $\phi$ is multiplicative, and that this function is a summatory function of $\phi$, so this function itself is multiplicative. We can decompose this into prime powers. So it suffices to show this for prime powers.

    Let $n = p^k$. Let
    \[f(n) = \sum_{d\mid n}\phi(d).\]
    Then we have
    \begin{align*}
        f(p^k) = \sum_{d\mid p^k} \phi(d) & = 1 + (p-1) + (p^2 - p) + \dots + (p^k-p^{k-1})
        \intertext{which is a telescoping sum which leaves}
                                          & = p^k
    \end{align*}
    which is as intended.
\end{proof}


\documentclass[../PHYS306Notes.tex]{subfiles}

\begin{document}
\section{Lecture 1}
\subsection{Lecture Notes - Lagrangian Mechanics Part 1}
\subsubsection{Newtonian Mechanics to Lagrangian Mechanics}
Recall the Newtonian formulation of classical mechanics; given the forces $\v{F}_1(t), \cdots \v{F}_n(t)$ on a particle, we can solve Newton's second law (a second order differential equation):
\[\sum_{i=1}^n \v{F}_i(t) = m \ddot{\v{r}}(t)\]
To obtain the trajectory $\v{r}(t)$, which is uniquely determined by the initial conditions $\v{r}(t_0)$ and $\dot{\v{r}}(t_0)$.

In this course, we will begin by looking at the Lagrangian formulation of classical mechanics. While this formulation contains no new physics compared to the Newtonian formulation, there are two distinct benefits:
\begin{enumerate}
    \item We can obtain EOM that do not depend on the coordinate system
    \item It is easier to treat constrained systems.
\end{enumerate}
\subsubsection{The Variational Principle Setup}
To do this, we will use a new approach, known as the \textbf{Variational principle}. To set this up, let us consider the trajectory (as well as some "wrong" paths between the same two endpoints) travelled by a particle:
\begin{center}
    \includegraphics[scale=0.4]{Lecture-1/L1-img1.png}
\end{center}
The trajectory from time $t_1$ to $t_2$ is parametrized by the generalized coordinate $q$. Examples of these are $x_1, y_1, \theta_1$.
\subsubsection{Generalized Coordinates}
For $N$ particles, a generalized coordinate $q_i$ depends on the positions of the $N$ particles:
\[q_i = q_i(\v{r}_1, \cdots, \v{r}_N)\]
In Cartesian coordinates, we have $3N$ coordinates but these are not necessarily independent. In general, we may have constraint functions $f_\alpha(\v{r}, \dot{\v{r}}, t) = 0$ where $\alpha = 1, \cdots, k$ (i.e. $k$ constraints). We then have $q_i$ independent coordinates, where $i = 1, \cdots n$ where $i = 3N - k$. In other words, for a system of $N$ particles, we have $3N - k$ independent/generalized coordinates. Generalized coordinates allow us to avoid worrying about the constrained parameters.
\subsubsection{Hamilton's Principle and the Lagrangian}
Let us consider assigning to each generalized coordinate $q(t)$ a number/value:
\[q(t) \mapsto S[q(t)] \in \RR\]
This is a functional (as denoted by the square brackets), as it takes in a function as an argument. Now, let us consider \textbf{Hamilton's principle}:
\begin{center}
    \textit{The actual path of a particle between times $t_1$ and $t_2$ is such that the line integral:
    \[ S[q] = \int_{t_1}^{t_2} \LL(q_1, \dot{q}_1, t) dt\]
    is stationary.}
\end{center}
Though the principle saysthe integral is stationary, often this correponds to a minimum (though not always). The function $\LL$ is defined as:
\[\LL = T - U\]
where $T$ is the kinetic energy and $U$ is the potential energy. This is called the \textbf{Lagrange Function} or \textbf{Lagrangian}. $S$ is called the \textbf{action}. Though we have not shown it explicitly, this Lagrange function gives the correct trajectory. Note that we have in essence replaced a second order ODE (Newton's second law) with an integral of $\LL$. Having two endpoints $t_1$ and $t_2$ is consistent with the two initial conditions for a second order ODE.
\end{document}


\documentclass{article}

\usepackage[margin=1in]{geometry}
\usepackage[utf8]{inputenc}
\usepackage{graphicx}
\usepackage{amsmath}
\usepackage{gensymb}

\newcommand{\norm}[1]{\lvert\lvert \, #1 \, \rvert\rvert}
\newcommand{\cross}{\times}
\newcommand{\pvec}[1]{\vec{#1}^{\,\prime}}

\title{Mobile Robot Kinematics for FTC}
\author{Ryan Brott}
\date{}

\begin{document}

\maketitle

\section{Introduction}

In FIRST Tech Challenge (FTC), kinematics are used all of the time---they're essential to robot operation. However, few teams  understand where the relationships come from or even acknowledge them explicitly, especially for holonomic drives. This paper intends to demystify the role and derivation of drive kinematics with a unified approach.

\section{Drive Kinematics}

This section systematically derives the forward and inverse kinematics for a variety of drives found in FTC (i.e., differential, mecanum, and swerve). The forward kinematics describe how the robot moves with specific wheel velocities. The inverse kinematics describe the opposite---how the wheels move during a given robot motion.

\subsection{Robots as Rigid Bodies}

Most FTC robots have a robust structure that holds all of the components together into one coherent system. When acted upon by a force, the robot does not materially deform, and its constituent parts remain essentially fixed relative to each other. Objects with this property are called rigid bodies, and their kinematics can be summarized with a single translational and angular velocity vector pair. The translational velocity describes the motion of the axis or tation, while the angular velocity magnitude points along the axis of rotation and has magnitude equal to the angular speed.\\
\\
Consider a point in the body with position $\vec{r}$. This position can be written as the sum $\vec{R} + \pvec{r}$ where $\vec{R}$ points from the origin to the axis of rotation and $\pvec{r}$ extends perpendicular to $\vec{\omega}$. Differentiation\footnote{The time derivative of a vector $\vec{v}$ in pure rotation described by $\vec{\omega}$ is $\frac{d\vec{v}}{dt} = \vec{\omega} \cross \vec{v}$.} gives the velocity relationship
\begin{align*}
    \frac{d\vec{r}}{dt} &= \frac{d\vec{R}}{dt} + \frac{d\pvec{r}}{dt} \\
    &= \vec{V} + \vec{\omega} \cross \pvec{r}.
\end{align*}
As FTC robots are generally confined to a plane, we have
\begin{align*}
    \vec{r} &= x \, \hat{\imath} + y \, \hat{\jmath}\\
    \vec{V} &= v_x \, \hat{\imath} + v_y \, \hat{\jmath}\\
    \vec{\omega} &= \omega \, \hat{k}.
\end{align*}
In this special case, the velocity is
\[
    \vec{v} = \frac{d\vec{r}}{dt} = (v_x - y \, \omega) \, \hat{\imath} + (v_y + x \, \omega) \, \hat{\jmath}.
\]
\\
Now for a given configuration $(v_x, v_y, \omega)$ it is possible to compute the two-dimensional velocity of any point on the robot.

\subsection{Differential Drives}

Perhaps the simplest practical mobile robot drive configuration is a differential drive with two normal wheels. These wheels have radius $R$ and are spaced $2l$ units apart. Let the center of rotation (i.e., the midpoint of the wheel positions) be the origin with the wheel axles laying parallel to the y-axis. With this setup, the wheel center positions are \((\pm l, 0)\).\\
\\
Using the results of the previous subsection, the desired velocities are
\begin{align*}
    \vec{v}_l &= (v_x - l \, \omega) \, \hat{\imath} + v_y \, \hat{\jmath}\\
    \vec{v}_r &= (v_x + l \, \omega) \, \hat{\imath} + v_y \, \hat{\jmath}.
\end{align*}
Furthermore, the orientation of the wheels gives us the tangential velocities in terms of the wheel angular velocities
\begin{align*}
    \pvec{v}_l &= \omega_l R \, \hat{\imath}\\
    \pvec{v}_r &= \omega_r R \, \hat{\imath}.
\end{align*}
For traction wheels, the corresponding velocities must \textit{exactly} match. This yields the equations
\begin{align*}
    v_x - l \, \omega &= \omega_l R \\
    v_x + l \, \omega &= \omega_r R \\
    v_y &= 0 \\
    v_y &= 0.
\end{align*}
From this we can solve to obtain the forward and inverse kinematics
\begin{equation*}
    \begin{aligned}[c]
        v_x &= \frac{R}{2}(\omega_l + \omega_r) \\
        v_y &= 0 \\
        \omega &= \frac{R}{2l}(\omega_r - \omega_l),
    \end{aligned}
    \qquad\qquad
    \begin{aligned}[c]
        \omega_l &= \frac{v_x - l \, \omega}{R} \\
        \omega_r &= \frac{v_x + l \, \omega}{R}.
    \end{aligned}
\end{equation*}
In matrix form, these relations are
\begin{equation*}
    \begin{bmatrix}
        v_x \\
        v_y \\
        \omega
    \end{bmatrix}
    =
    \frac{R}{2}
    \begin{bmatrix}
        1 & 1 \\
        0 & 0 \\
        -\frac{1}{l} & \frac{1}{l}
    \end{bmatrix}
    \begin{bmatrix}
        \omega_l \\
        \omega_r
    \end{bmatrix},
    \qquad\qquad
    \begin{bmatrix}
        \omega_l \\
        \omega_r
    \end{bmatrix}
    =
    \frac{1}{R}
    \begin{bmatrix}
        1 & 0 & -l\\
        1 & 0 & l
    \end{bmatrix}
    \begin{bmatrix}
        v_x \\
        v_y \\
        \omega
    \end{bmatrix}.
\end{equation*}
These same kinematics can be applied to differential drives with multiple wheels per side. However, you'll quickly realize that it is no longer possible to satisfy all the constraints. To rotate properly, the wheels must overcome static friction and slide parallel to the wheel axis. This slippage can be mathematically accounted for by treating this scrub direction as a passive degree of freedom (like an omni wheel) as we'll see in the next section. With this assumption, the kinematics are the same as above with all motors on each side sharing the same velocity.

\subsection{Mecanum Drives}

Mecanum drives consist of four mecanum wheels with rollers at $45\degree$ angles with respect to the axle positioned in a rectangle. Looking from above, the main diagonal wheels have counterclockwise-rotated rollers, while the alternate diagonal wheels have clockwise-rotated rollers. The careful combination of these two different directions enables omnidirectional movement. The track width (distance between opposing wheels) is $2l$ and the wheelbase (distance between adjacent wheels) is $2b$. Place the center of rotation at the origin as before with wheel center positions \((\pm l, \pm b)\).\\
\\
Using the same procedure as before, we obtain
\begin{align*}
    \vec{v}_{fl} &= (v_x - l \, \omega) \, \hat{\imath} + (v_y + b \, \omega) \, \hat{\jmath}\\
    \vec{v}_{bl} &= (v_x - l \, \omega) \, \hat{\imath} + (v_y - b \, \omega) \, \hat{\jmath}\\
    \vec{v}_{br} &= (v_x + l \, \omega) \, \hat{\imath} + (v_y - b \, \omega) \, \hat{\jmath}\\
    \vec{v}_{fr} &= (v_x + l \, \omega) \, \hat{\imath} + (v_y + b \, \omega) \, \hat{\jmath}.
\end{align*}
Computing the tangential velocities from the angular velocities is bit trickier for mecanum wheels. The tangential force acts $45\degree$ from the roller, resulting in a speed reduction of $\sqrt{2}$. The direction of the velocity points perpendicular to the roller axis (viewed from above). Thus, the velocities are
\begin{align*}
    \pvec{v}_{fl} &= \frac{R \, \omega_{fl}}{2} (\hat{\imath} - \hat{\jmath})\\
    \pvec{v}_{bl} &= \frac{R \, \omega_{bl}}{2} (\hat{\imath} + \hat{\jmath})\\
    \pvec{v}_{br} &= \frac{R \, \omega_{br}}{2} (\hat{\imath} - \hat{\jmath})\\
    \pvec{v}_{fr} &= \frac{R \, \omega_{fr}}{2} (\hat{\imath} + \hat{\jmath}).
\end{align*}
The vectors can't be directly equated for non-square configurations due to the passive motion of the rollers. Instead we equate the portion of each vector in the tangential velocity direction. For example, the front left constraint is
\[
    \vec{v}_{fl} \cdot \pvec{v}_{fl} = \pvec{v}_{fl} \cdot \pvec{v}_{fl}.
\]
The final system is
\begin{align*}
    \frac{R \, \omega_{fl}}{2}(v_x - l \, \omega - v_y - b \, \omega) &= \frac{R^2\omega_{fl}^2}{2}\\
    \frac{R \, \omega_{bl}}{2}(v_x - l \, \omega + v_y - b \, \omega) &= \frac{R^2\omega_{bl}^2}{2}\\
    \frac{R \, \omega_{br}}{2}(v_x + l \, \omega - v_y + b \, \omega) &= \frac{R^2\omega_{br}^2}{2}\\
    \frac{R \, \omega_{fr}}{2}(v_x + l \, \omega + v_y + b \, \omega) &= \frac{R^2\omega_{fr}^2}{2}.
\end{align*}
which yields the matrix kinematics
\begin{equation*}
    \begin{bmatrix}
        v_x \\
        v_y \\
        \omega
    \end{bmatrix}
    =
    \frac{R}{4}
    \begin{bmatrix}
        1 & 1 & 1 & 1 \\
        -1 & 1 & -1 & 1 \\
        -\frac{1}{l+b} & -\frac{1}{l+b} & \frac{1}{l+b} & \frac{1}{l+b}
    \end{bmatrix}
    \begin{bmatrix}
        \omega_{lf} \\
        \omega_{lb} \\
        \omega_{rb} \\
        \omega_{rf}
    \end{bmatrix}
    \qquad\qquad
    \begin{bmatrix}
        \omega_{lf} \\
        \omega_{lb} \\
        \omega_{rb} \\
        \omega_{rf}
    \end{bmatrix}
    =
    \frac{1}{R}
    \begin{bmatrix}
        1 & -1 & -(l + b) \\
        1 & 1 & -(l + b) \\
        1 & -1 & (l + b) \\
        1 & 1 & (l + b)
    \end{bmatrix}
    \begin{bmatrix}
        v_x \\
        v_y \\
        \omega
    \end{bmatrix}
\end{equation*}

\subsection{Swerve Drives}
Unlike the previous drives, swerve drives have actuated wheel directions. Nevertheless, swerve kinematics can still be derived within the same framework. Let each wheel have coordinates $(x_i, y_i)$, orientation $\phi_i$, tangential velocity $\vec{v}_i$, and angular velocity $\omega_i$.\\
\\
For simplicity, we will assume the swerve wheels do not slip and directly equate the tangential velocities: $v_x - y_i \, \omega = R \, \omega_i \operatorname{cos} \phi_i$ and $v_y + x_i \, \omega = R \, \omega_i \operatorname{sin} \phi_i$. This gives the following inverse kinematics\footnote{$\operatorname{atan2}$ is the standard two-argument arctangent function available in most programming environments}:
\begin{align*}
    \phi_i &= \operatorname{atan2} \big( v_y + x_i \, \omega, \; v_x - y_i \, \omega \big)\\
    \omega_i &= \frac{1}{R} \sqrt{(v_x - y_i \, \omega)^2 + (v_y + x_i \, \omega)^2}
\end{align*}

Forward kinematics for a collection of multiple wheels can be computed by solving the corresponding overdetermined system.

\section{Odometry}
The previous section focused on finding the relationship between configuration and wheel velocities. This by itself is sufficient to employ the inverse kinematics for sending the appropriate control signals. However, the forward kinematics cannot be directly used for odometry. To localize the robot, the local velocities must be integrated into positions and transformed into the global frame.

\subsection{Constant Velocity Odometry}
For simplicity, this method assumes constant translational and rotational velocity over each measurement period. In practice, this is a good assumption so long as measurements are frequent enough (additionally, estimating acceleration or other higher order derivatives robustly from wheel position data is nontrivial).\\
\\
Without loss of generality, we take the robot heading $\theta$ to be $0$ initially. A measurement is then taken $\Delta t$ time later. During this period the robot's global velocity is
$$
\begin{bmatrix}
    \dot{x}_G \\
    \dot{y}_G \\
    \dot{\theta}_G
\end{bmatrix}
=
\begin{bmatrix}
    \operatorname{cos} \theta & -\operatorname{sin} \theta & 0 \\
    \operatorname{sin} \theta & \operatorname{cos} \theta & 0 \\
    0 & 0 & 1
\end{bmatrix}
\begin{bmatrix}
    \dot{x}_R \\
    \dot{y}_R \\
    \dot{\theta}_R
\end{bmatrix}.
$$
Integrating this over time gives
$$
\begin{bmatrix}
    \Delta x_G \\
    \Delta y_G \\
    \Delta \theta_G
\end{bmatrix}
=
\begin{bmatrix}
    \frac{\operatorname{sin} \Delta \theta_R}{\Delta \theta_R} & -\frac{1 - \operatorname{cos} \Delta \theta_R}{\Delta \theta_R} & 0 \\[4pt]
    \frac{1 - \operatorname{cos} \Delta \theta_R}{\Delta \theta_R} & \frac{\operatorname{sin} \Delta \theta_R}{\Delta \theta_R} & 0 \\[4pt]
    0 & 0 & 1
\end{bmatrix}
\begin{bmatrix}
    \Delta x_R \\
    \Delta y_R \\
    \Delta \theta_R
\end{bmatrix}.
$$
This can then be rotated and added to the previous estimate to complete the update
$$
\begin{bmatrix}
    x_{G,\,t+1} \\
    y_{G,\,t+1} \\
    \theta_{G,\,t+1}
\end{bmatrix}
=
\begin{bmatrix}
    x_{G,\,t} \\
    y_{G,\,t} \\
    \theta_{G,\,t}
\end{bmatrix}
+
\begin{bmatrix}
    \operatorname{cos} \theta_{G,\,t} & -\operatorname{sin} \theta_{G,\,t} & 0 \\
    \operatorname{sin} \theta_{G,\,t} & \operatorname{cos} \theta_{G,\,t} & 0 \\
    0 & 0 & 1
\end{bmatrix}
\begin{bmatrix}
    \frac{\operatorname{sin} \Delta \theta_R}{\Delta \theta_R} & -\frac{1 - \operatorname{cos} \Delta \theta_R}{\Delta \theta_R} & 0 \\[4pt]
    \frac{1 - \operatorname{cos} \Delta \theta_R}{\Delta \theta_R} & \frac{\operatorname{sin} \Delta \theta_R}{\Delta \theta_R} & 0 \\[4pt]
    0 & 0 & 1
\end{bmatrix}
\begin{bmatrix}
    \Delta x_R \\
    \Delta y_R \\
    \Delta \theta_R
\end{bmatrix}.
$$

\subsection{Tracking Wheels}
Tracking wheels are passive omni wheels intended solely for odometry. Each wheel has an arbitrary position $(x_i, y_i)$ and orientation $\phi_i$ in the robot frame. Similar to swerve, the tangential velocities of each wheel are $\vec{v}_i = (v_x - y_i \, \omega) \, \hat{\imath} + (v_y + x_i \, \omega) \, \hat{\jmath}$ and $\pvec{v}_i = R \, \omega_i (\operatorname{cos} \phi_i \, \hat{\imath} + \operatorname{sin} \phi_i \, \hat{\jmath})$.
$$
    \omega_i = \frac{1}{R}\Big[v_x \operatorname{cos} \phi_i + v_y \operatorname{sin} \phi_i + \omega \, (x_i \operatorname{sin} \phi_i - y_i \operatorname{cos} \phi_i)\Big]
$$
If placed properly, three tracking wheels are sufficient to determine the $(v_x, v_y, \omega)$ configuration:
$$
\begin{bmatrix}
    \omega_1\\
    \omega_2\\
    \omega_3
\end{bmatrix}
=
\frac{1}{R}
\begin{bmatrix}
    \operatorname{cos} \phi_1 & \operatorname{sin} \phi_1 & x_1 \operatorname{sin} \phi_1 - y_1 \operatorname{cos} \phi_1 \\
    \operatorname{cos} \phi_2 & \operatorname{sin} \phi_2 & x_2 \operatorname{sin} \phi_2 - y_2 \operatorname{cos} \phi_2 \\
    \operatorname{cos} \phi_3 & \operatorname{sin} \phi_3 & x_3 \operatorname{sin} \phi_3 - y_3 \operatorname{cos} \phi_3 \\
\end{bmatrix}
\begin{bmatrix}
    v_x\\
    v_y\\
    \omega
\end{bmatrix}
$$
$$
\begin{bmatrix}
    \Delta \theta_1\\
    \Delta \theta_2\\
    \Delta \theta_3
\end{bmatrix}
=
\frac{1}{R}
\begin{bmatrix}
    \operatorname{cos} \phi_1 & \operatorname{sin} \phi_1 & x_1 \operatorname{sin} \phi_1 - y_1 \operatorname{cos} \phi_1 \\
    \operatorname{cos} \phi_2 & \operatorname{sin} \phi_2 & x_2 \operatorname{sin} \phi_2 - y_2 \operatorname{cos} \phi_2 \\
    \operatorname{cos} \phi_3 & \operatorname{sin} \phi_3 & x_3 \operatorname{sin} \phi_3 - y_3 \operatorname{cos} \phi_3 \\
\end{bmatrix}
\begin{bmatrix}
    \Delta x_R\\
    \Delta y_R\\
    \Delta \theta_R
\end{bmatrix}
$$
For example, take the common configuration of two tracking wheels placed parallel to the robot's $x$-axis at $(0,\pm y_0)$ and a third placed parallel to the robot's $y$-axis at $(x_0,0)$. For simplicity, all wheels point in the positive direction of their corresponding axis. This gives the following matrices:
\begin{equation*}
    \begin{bmatrix}
        \Delta \theta_1\\
        \Delta \theta_2\\
        \Delta \theta_3
    \end{bmatrix}
    =
    \frac{1}{R}
    \begin{bmatrix}
        1 & 0 & -y_0 \\
        1 & 0 & y_0 \\
        0 & 1 & x_0 \\
    \end{bmatrix}
    \begin{bmatrix}
        \Delta x_R\\
        \Delta y_R\\
        \Delta \theta_R
    \end{bmatrix}
\end{equation*}
\begin{equation*}
    \begin{bmatrix}
        \Delta x_R\\
        \Delta y_R\\
        \Delta \theta_R
    \end{bmatrix}
    =
    \frac{R}{2y_0}
    \begin{bmatrix}
        y_0 & y_0 & 0\\
        x_0 & -x_0 & 2y_0\\
        -1 & 1 & 0\\
    \end{bmatrix}
    \begin{bmatrix}
        \Delta \theta_1\\
        \Delta \theta_2\\
        \Delta \theta_3
    \end{bmatrix}
    \quad\leftrightarrow\quad
    \begin{aligned}[c]
        \Delta x_R &= \frac{R}{2}(\Delta \theta_1 + \Delta \theta_2)\\
        \Delta y_R &= R\bigg[\frac{x_0}{2y_0}(\Delta \theta_1 - \Delta \theta_2) + \Delta \theta_3\bigg]\\
        \Delta \theta_R &= \frac{R}{2y_0}(\Delta \theta_2 - \Delta \theta_1)
    \end{aligned}
\end{equation*}
Two tracking wheels and a heading sensor is also sufficient for localization.
$$
\begin{bmatrix}
    \Delta \theta_1\\
    \Delta \theta_2\\
    \Delta \theta_3
\end{bmatrix}
=
\frac{1}{R}
\begin{bmatrix}
    \operatorname{cos} \phi_1 & \operatorname{sin} \phi_1 & x_1 \operatorname{sin} \phi_1 - y_1 \operatorname{cos} \phi_1 \\
    \operatorname{cos} \phi_2 & \operatorname{sin} \phi_2 & x_2 \operatorname{sin} \phi_2 - y_2 \operatorname{cos} \phi_2 \\
    0 & 0 & R\\
\end{bmatrix}
\begin{bmatrix}
    \Delta x_R\\
    \Delta y_R\\
    \Delta \theta_R
\end{bmatrix}
$$

\end{document}


\lab{Value Function Iteration}{Value Function Iteration}
\newcommand\ve{\varepsilon}
\objective{This section teaches the fundamentals of Dynamic Programming using value function iteration.}

%Often it is of interest to optimize decision making in some sequential process.  For example, an oil company may need to decide
%how much oil to excavate and sell each month as prices change, a person entering retirement may need to decide how much of their
%savings to spend each year, or a model of economic growth may require a decision about how much to invest in capital versus how
%much to spend each year.  In this lab we will formulate a general dynamic optimization problem.  We will explore techniques for
%solving such a problem with both finite and infinite time horizons.

Previously we have optimizated a single decision or question.
In this lab we will optimize a sequence of decisions.
The techniques to solve these problems are called dynamic optimization.

Dynamic optimization provides different answers than optimization techniques we have studied previously.
For example, an oil company might want to know how much oil to excavate in one day in order to maximize profit.
If we only consider today, then the answer is to excavate and sell as much as possible.
If each day is considered in isolation, then the strategy to optimize profit is to maximize excavation in order to maximize profits.
However, in reality oil prices change from day to day and as supply increases or decreases, and so maximizing excavation may in fact lead to less profit.
On the other hand, if the oil company considers how production on one day will affect subsequent decisions, they may be able to maximize their profits.
In this lab we will explore techniques for solving such a problem.

\section*{The Cake Eating Problem}

Rather than maximizing oil profits, we will focus on solving a general problem that can be applied in many areas called the cake eating problem.
Given a cake of a certain size, how do we eat it to maximize our enjoyment (in other words, our utility) over time?
Some people may prefer to eat all of their cake at once, and not save any for later.
Others may prefer to only eat a little bit at a time.
These preferences will be expressed precisely using a utility function.
Our task is to find an optimal strategy given an increasing utility function, $u$.
Precisely, given a cake of size $W_0$ and some amount of consumption $c_0 \in [0, W_0]$, the utility gained is given by
\[
u(c_0).
\]

If we want to maximize utility over one time period, we will consume the entire cake.  How do we maximize utility over several days?

\subsection*{Discount Factors}

A person or firm will typically have a preference for saving or consuming.
For example, a dollar today can be invested and yield interest, whereas a dollar received next year will not include the accrued interest.
In this lab, cake now will yield more utility than cake in the future.
We can model this by multiplying future utility by a discount factor $\beta \in (0,1)$.
For example, if we were to consume $c_0$ cake at time $0$ and $c_1$ cake at time $1$, then the total utility gained will be
\[
u(c_0) + \beta u(c_1).
\]

\subsection*{The Optimization Problem}

If we are to consume a cake of size $W_0$ over $T$ time periods, then we can represent our consumption at each step as a vector
\[
(c_0, c_1, ... , c_T)
\]
where
\[
\sum_{i=1}^T c_i = W_0.
\]

We will call such a vector a policy.  Thus, our optimization problem is to

\begin{align*}
\mbox{maximize }  & \sum_{t=0}^T \beta^t u(c_t) \\
\mbox{subject to } & \sum_{t=0}^T c_t = 1, c_t \geq 0.
\end{align*}

\begin{problem}

It might seem obvious what sort of policy will yield the most utility, but the truth may surprise you.
See Figure \ref{fig:diff_pols} for some examples.
Write a function called \li{graph_policy} will accepts
\begin{itemize}
\item a policy,
\item a utility function, and
\item a discount factor;
\end{itemize}
and returns the total utility gained with the policy input.
Also display a plot of the total utility gained over time.
Ensure that the policy that the user passes in sums to 1.
Otherwise, raise an \li{InputError}.

%\begin{figure}
%\begin{subfigure}{.5\textwidth}
%    \includegraphics[width=\textwidth]{fixed_time.pdf}
%\end{subfigure}
%\begin{subfigure}{.5\textwidth}
%    \includegraphics[width=\textwidth]{fixed_w.pdf}
%\end{subfigure}
%\caption{Slices of the finite horizon value function for fixed values of $t$ and $W$, respectively.}
%\label{fig:valueslices}
%\end{figure}
\end{problem}
\begin{figure}
\includegraphics[width=\textwidth]{diff_policies.pdf}
\caption{Plots for various policies with a square root utility function.  Policy 1 eats all the cake in the first step while policy 2 eats all the cake in the last step.  Their difference in utility demonstrate the effect of the discount factor on waiting to eat.  Policy 3 eats the same amount of cake at each step, while policy 4 starts eating a lot of cake but eats less and less as time goes on and the cake runs out.}
\label{fig:diff_pols}
\end{figure}

\section*{The Value Function}

We have posed the cake eating problem as an optimization problem, but we do not have a good way to solve it yet.
To that end, we introduce the value function.
The value function  $V(a,b,W)$ is the maximal value of

\begin{align*}
\mbox{maximize } & \sum_{t=a}^b \beta^tu(c_t) \\
\mbox{subject to } & \sum_{t=a}^b c_t = W, c_t \geq 0
\end{align*}

for $t = a, a+1, ... b$.
In other words, it gives the utility gained from following an optimal policy from time $a$ to time $b$.
%The Value function gives how much utility that a policy from time $a$ to $b$ yields given a cake of size $W$.
So if we start with a cake of size $W_0$ at time $0$ and have half the cake left at time $t$, then $V(t, T, \frac{W_0}{2})$ gives how much utility we will gain by proceeding optimally from $t$.

Thus, if we know what the optimal policy in the future is, we can figure out the optimal policy for the present, since

\[
V(t, T, W_{t}) = \max_{W_{t+1}} u(W_{t} - W_{t+1}) + \beta V(t+1, T, W_{t+1}).
\]
Recall that we know what the utility function $u$ is, as well as $\beta$.
Also, note that our consumption at time $t$ is given by the amount of cake at the beginning of time $t$ minus the amount of cake at the beginning of time $t+1$, or $W_t - W_{t+1}$.

In other words, if we know how to maximize utility in the future, we only need to maximize utility in the now in order to maintain optimality.
So if we know what $V(t+1, T, W_{t+1})$ yields for all possible $W_{t+1}$, then we choose the best $W_{t+1}$ to figure out what the value function yields in the present.
Thus, if we can evaluate the value function, we can also determine the optimal policy.

\subsection*{Solving the Optimization Problem}

\begin{problem}

In this problem we will write a helper function that will help us choose the optimal policy.
We will assume our discritized cake has volume $1$ and an integer $N$ equaly-sized pieces.
Write a method that accepts $N$ and a utility function.
Create a vector whose entries correspond to possible amounts of cake given the discretization.
For example, if we wanted to consider a cake split into 4 pieces, the vector should be
\[
w = (0, 0.25, 0.5, 0.75, 1.0)
\]
Return a matrix whose $ij^{th}$ entry is the amount of utility gained by starting with $i$ cake pieces and ending with $j$ pieces.
In other words, the $ij^{th}$ entry should be $u(w_i - w_j)$.
Let the default utility function be the square root function.
Set impossible situations to 0 (i.e., eating more than you have).
We call the resulting matrix the consumption matrix.

\end{problem}

How can we find the value function if we do not know the optimal policy?
So far in our problem, we do not know how to start eating in order to maximize utility.
However, we know that we do not want to have any cake left over when time runs out.
Thus, the amount of utility we gain from having $W_{T-1}$ cake at time $T-1$ is given by $u(W_{T-1})$, and so we know $V(T-1, T, W_{T-1})$ given a choice of $W_{T-1}$.
For example, if we had a cake split into 4 pieces and wanted to eat it in $T=4$ time intervals, we could express what we know about the value function from time $T$ to $T+1$ as

\[
\begin{bmatrix}
* & * & * & * & u(1) & 0 \\
* & * & * & * & u(0.75) & 0 \\
* & * & * & * & u(0.5) & 0 \\
* & * & * & * & u(0.25) & 0 \\
* & * & * & * & :w
u(0) & 0 \\
\end{bmatrix}
\]

where the $ij^{th}$ entry gives the value of having $w_j$ cake at time $j$ and $*$ represents unknown values.
This is called the value function matrix, and note that it is dimension $N \times T+2$, since we want columns for $t = 0$ and $t = T+1$.
The last column is there to show that after time $4$, there is no value in having any cake.

\begin{problem}
Write a function that accepts
\begin{itemize}
\item $T$ (the number of time steps),
\item $N$ (the number of pieces into which we split our cake),
\item a discount factor $\beta$, and
\item a utility function;
\end{itemize}
and returns the value function matrix with the last two columns filled in.
Let the default utility function be the square root function.
\end{problem}

As we mentioned in the previous section, if we know what the value function evaluates to in the future, we can evaluate the value function in the present.
If we are at time $t$, and we know what the value function evaluates to at time $t+1$ (given by the appropriate column of the value function matrix we have built), then we can calculate
\[
V_t(W_t) = u(W_{t} - W_{t+1}) + \beta V_{t+1}(W_{t+1})
\]
for all possible values of $W_t$ and $W_{t+1}$.
This will produce an $N \times N$ matrix giving the value of all possible values of starting with a cake of size $W_t$ and finishing with size $W_{t+1}$.
We call this matrix the current value matrix.
%If we choose the maximal value of each row, then we know the optimal amount of cake to end up with a the end of time $t$, and thus how much cake to eat.
The largest entry of each row then is the optimal value that the value function can attain at this step, given that we start with the amount of cake that corresponds to the row.
Thus, the maximal values of each row become the column of the value function matrix at time $t$ and if we know one column of the value matrix, we may iterate backwards to fill the rest of it in.
Note that we can start this process with just the column of zeros that we know occupies the last column of the value function matrix.

\begin{problem}
Modify the solution to the previous problem to determine the entire value function matrix.
Initialize the matrix as zeros, and starting from the next to last column
\begin{itemize}
\item calculate the current value matrix,
\item find the largest value in each row, and
\item fill in the column with these values.
\end{itemize}
\end{problem}

Now that the value function matrix is filled in we could say that the optimization problem is solved in some sense.
However, it is not immediately apparent how to make decisions given the value function matrix only.
What we really want is the optimal policy.
It turns out that we already calculated it in the previous problem - we just didn't keep track of it.
Recall that we calculated a current value matrix for each column of the value matrix, and that the $ij^{th}$ entry of that matrix was the value of starting with $i$ pieces of cake and finishing with $j$ pieces of cake.
Thus, the policy for starting a time step with $i$ pieces of cake is $w_i - w_j$ (the vector $w = (w_1, ..., w_N)$ is the equal partition of the cake), where $j$ is the index of w that corresponds to the maximum value of row $i$ in the current value matrix.
We can use this information to fill out an $N \times T+1$ policy matrix at the same time that we fill out the value matrix (T+1 because we inlcude $t=0$).
The $ij^{th}$ entry of the policy matrix tells us how much cake to eat if we start with $w_i$ cake at time $j$.

\begin{problem}
Modify the solution to the previous problem to determine the policy matrix.
Initialize the matrix as zeros and fill it in starting from the last column at the same time that you determining the value function matrix.
\end{problem}

\begin{problem}
The $ij^{th}$ entry of the policy matrix tells us how much cake to eat at time $j$ if we start with $i$ pieces.
Use this information to find the optimal policy for starting with a cake of size 1 split into $N$ pieces, and then plot the policy.
\end{problem}

A summary of the arrays generated in this lab is given below, in the order that they were generated in the lab:

Consumption matrix: Equal to $u(u_i - w_j)$, the utility gained when you start with $i$ pieces and end with $j$ pieces.

Value Function matrix: How valuable is it to have $w_i$ pieces at time $j$.

Current Value matrix: How valuable is each possible decision at time $t$.

Policy matrix: The amount of cake we decide to eat at time $t$.

\begin{comment}
%
% This is the old version of the lab.  I am keeping it in for now.
% If we wish to expand this lab then this material may help write
% more sections, i.e. infinite horizon porblems.
%

\section*{The Sequential Problem, Finite Horizon}
Suppose there are time periods $t=0,1,\ldots, T$ and at each time period we take an action $c_t$. Furthermore, at the beginning
of each time period $t$ we are in some state $W_t$.  In many cases $W_t$ might represent an available resource, such as money.
At each time we receive some reward, $u(W_t,c_t)$, for taking action $c_t$ given state $W_t$.  We assume that rewards are worth
more now than later. We let $\beta\in (0,1)$ represent what is called the discount factor, which gives the ratio of preference for
rewards today versus rewards tomorrow.  For example, receiving a dollar today is preferable to receiving a dollar in a year
because taking a dollar today and putting it into a savings account results in having more than a dollar in a year.  Lastly, over
time our state variable $W_t$ changes according to some rule depending on the previous state and our actions,
\begin{equation}
\label{motion}
W_{t+1} = g(W_t,c_t).
\end{equation}
Equation \eqref{motion} is sometimes referred to as the law of motion, as it describes how we move from state to state.
Mathematically such a problem can be represented as follows:
\begin{equation*}
\text{maximize} \sum_{t=0}^T \beta^t u(W_t,c_t) \quad \text{s.t.} \quad W_{t+1} = g(W_t,c_t)
\end{equation*}
where our initial state, $W_0$, is given.  There may also be restrictions on our choices $c_t$.  For example, in many applications
the state $W_t$ represents the amount of some resource available, and $c_t$ represents the amount we use up in time period $t$.
 In this case we would require $c_t \in [0,W_t]$.

For simplicity, lets assume that $u$ is a function of $c_t$ only (this is often, though not always, the case in practice).
 First let's consider the case that $T=0$.  So we maximize
\begin{equation}\label{1perprob}
u(c_0)
\end{equation}
over $c_0 \in [0,W_0]$.  In most cases $u$ is increasing, which we will assume here.  In this case it will be optimal to choose
 the largest value of $c_0$ possible, that is, $c_0 = W_0$.  Thinking of $W_0$ as our available resources, this simply means
 that if we don't have future periods to consider, we will use all of it.

In fact, this is always true in the last period.  In a problem with $T$ periods we know that we will use all of our resources
remaining in period $T$.  Consider the two period problem:
\begin{equation}
\label{2perprob}
\max \, \{u(c_0) + \beta u(c_1)\}
\end{equation}
where $c_0 \in [0,W_0]$, $c_1 \in [0,W_1]$ and $W_1 = g(W_0,c_0)$.  We know that in the last period we will use all of our
remaining resources, so $c_1 = W_1=g(W_0,c_0)$.  Substituting gives
\begin{equation}
\label{2perproba}
\max \, \{u(c_0) + \beta u(g(W_0,c_0))\}.
\end{equation}
Now we need only determine $c_0$.  Taking the derivative of \eqref{2perproba} with respect to $c_0$ and setting equal to zero
gives the first order condition
\begin{equation}
\label{FOC}
u'(c_0) = -\beta u'(g(W_0,c_0))g_c(W_0,c_0)
\end{equation}
where $g_c$ is the partial derivative of $g$ with respect to $c_0$.

Given a specific form for $u$ we could solve for $W_1$ and obtain the optimal solution.  In fact, we can solve a problem of
any length $T$ in this manner, by starting at the last time period and working backward.  We know that $W_{T+1} = 0$.  Working
backward in time we obtain an equation at each time step $t<T$ by taking the derivative with respect to $c_t$ and setting the equation equal to zero.  This process is called backward induction.  The equations at each time step, such as equation \eqref{FOC}, are sometimes called the inter-temporal Euler equations.  These equations, along with $c_T = W_T$, make $T+1$ equations to go with our $T+1$ unknowns $\{c_0,c_1,c_2,\ldots,c_T\}$, where we can use the law of motion \eqref{motion} to relate the $c_t$ and $W_t$

\section*{The Recursive Problem, Finite Horizon}
Approaching the problem sequentially like this can be somewhat messy.  The dynamic programming approach we consider now is more
easily adaptable to many situations.  The key to the dynamic programming approach is to define our optimization problem in terms
 of subproblems.  Notice that if we are in time period $t$, we face a problem of exactly the same form as the problem at time
 $0$.  We are in some state $W_t$, and want to maximize the sum from $t$ to $T$.  With this idea in mind, we define a function
  $V_t(W_t)$ called the value function.  The function $V_t$ gives the value of entering time $t$ in state $W_t$ and making
  optimal decisions moving forward.  So
\begin{equation*}
V_{t-1}(W_{t-1}) = \max_{c_t} \left\{u(c_{t-1}) + \beta V_t(g(W_{t-1},c_{t-1}))\right\}.
\end{equation*}
This is called the Bellman Equation.  The key to this formulation is that we decide what to do in period $t-1$ with the
assumption that our actions in the remaining periods will be optimal.  This is called the principal of optimality.

Let us consider a specific example from economics called The Cake Eating Problem.  Suppose $W_t$ represents the amount of
cake available at time $t$.  At each time period we can choose how much to consume.  What we eat, $c_t$, gives us a reward.
What we save, $W_t-c_t$, does not give us a reward (until it is eaten in a later period).  The law of motion \eqref{motion}
becomes
\begin{equation}\label{LOM_EX}
W_{t+1} = g(W_t,c_t) = W_t-c_t
\end{equation}

Now we have completely defined the problem.  The Bellman Equation is
\begin{equation*}
V_{t-1}(W_{t-1}) = \max_{c_t} \left\{u(c_{t-1}) + \beta V_t(W_{t-1}-c_{t-1})\right\}.
\end{equation*}
Notice that by the law of motion, each $c_t$ is determined by $W_t$ and $W_{t+1}$.
In fact, rearranging \eqref{LOM_EX}, we have
\begin{equation*}
c_t = W_t - W_{t+1}.
\end{equation*}
We can therefore rewrite the value function as
\begin{equation}
V_{t-1}(W_{t-1}) = \max_{W_t} \left\{u(W_{t-1} - W_{t}) + \beta V_t(W_t)\right\}.
\label{cake_valfn}
\end{equation}

We see that determining the optimal actions $c_t$ is equivalent to determining the optimal states $W_{t}$ in the
above formulation. The solution to this
problem is often called a \emph{policy function}.  A policy function determines an action based on the current
state.  Denoting the policy function by $\psi$, this can be written as
\begin{equation*}
W_{t+1}=\psi_t \left(W_t\right).
\end{equation*}
The policy function gives the optimal amount of cake to leave for the next period (equivalent to the amount of consumption) given
the amount of cake at the start of the period.  In other words, it determines the choice of $W_t$ that satisfies the $\max$
condition in \eqref{cake_valfn}.

As before, we know that in the last time period we should not save anything.  So $V_{T+1}(W_{T+1}) = 0$, i.e. there is
no value in leaving wealth for period $T+1$. Stated in another way, our action at time $T$ should be to eat all of the
remaining cake $W_T$, so $W_{T+1} = \psi_T(W_T) = 0$.
Plugging this result into the Bellman Equation gives us $V_T(W_T) = u(W_T)$.
Now consider the value function equation for period $T-1$:
\begin{align*}
V_{T-1}(W_{T-1}) &= \max_{W_T} \left\{u(W_{T-1} - W_T) + \beta V_T(W_T)\right\} \\
                 &= \max_{W_T} \left\{u(W_{T-1} - W_T) + \beta u(W_T)\right\}.
\end{align*}
We can determine this value by optimizing over $W_T$, where $0 \leq W_T \leq W_{T-1}$.
Continuing backwards in this manner leads us to the solution of the original problem.

\begin{problem}
\label{prob:cake_prob}
We recommend reading the entire problem before beginning to work on it, as many questions may be addressed further on.  This applies to the other problems in this lab as well.

Follow the steps below to solve the problem described above.  Take $u(c_t) = \sqrt{c_t}$.
You will write a function called \li{eatCake} that takes parameters $\beta$ (the discount factor),
$N$ (the number of discrete cake values to consider), $W_{max}$ (the original size of the cake,
set to the default value of $1$), a keyword argument \li{finite} (set to default value \li{True}),
a keyword argument  $T$ (the number of time periods, set to
default value \li{None}), and a keyword argument \li{plot}, which indicates whether or not
to plot the computed results. The function should return arrays representing the value function and the
policy function (we describe how to compute these in the following steps).
\begin{enumerate}
\item Approximate the continuum of possible cake sizes by creating an array
of evenly-spaced values that range from to 0 to $W_{max}$ inclusive.
Let the number of possible cake values be given by $N$. In Python, this can be accomplished easily by
using the \li{linspace} function in NumPy. You should obtain an array (call it $w$) of the form
\[
w = (w_1, w_2, \ldots, w_N),
\]
where $w_1 = 0$ and $w_N = W_{max}$.

\item Note that in order to compute the value function, we need $u(W_{t-1}-W_t)$.
We will pre-compute all such possible values and store them in an array, as follows:
Create an $N$ by $N$ matrix that contains
all possible values of $W_{t-1} - W_t$ (where $W_{t-1}$ corresponds to rows and $W_{t}$ to columns).  Make sure that
$c_t \geq 0$ is satisfied by replacing negative entries in the matrix with zeros.  Then take the square root to get a matrix
of $u(W_{t-1}-W_t)$.  To make sure we do not choose $W_{t-1} - W_t < 0$ when maximizing, replace the corresponding entries of
the $u(W_{t-1}-W_t)$ matrix with a large negative number (e.g. $-10^{10}$). You should end up with a matrix whose
$(i,j)$-th entry is equal to $\sqrt{w_i - w_j}$ when $i \geq j$, and is equal to $-10^{10}$ when $i < j$.

\item Next, create an $N$ by $T+2$ (corresponding to $t=0,1,\ldots, T+1$) matrix representing the value function for a given
time $t$ and state $W_t$.  We can initialize it with zeros and begin filling in the columns starting with the last (which we
know is zeros), as explained below.

\item Now we are ready to iterate backward and compute the value function for each time period.  To find $V_T$, we first compute
$u(W_T - W_{T+1}) + \beta V_{T+1}(W_{T+1})$ for all values of $W_{T}$ and $W_{T+1}$.  This will result in an $N$ by $N$ matrix
where the rows correspond to values of $W_{T}$ and the columns correspond to values of $W_{T+1}$.
%Note that to compute this
%we need a matrix representing $\beta V_{T+1}(W_{T+1})$.  Because this quantity does not depend on $W_T$, its rows should be
%equal.  To do this, we want to take $\beta V_{T+1}(W_{T+1})$ as a row vector and stack this vector to create a matrix with
%equal rows.  There are multiple ways to do this.  One is the \li{np.repeat} function.  For example, if \li{b} is a row vector
%it could be used like the following.
%\begin{lstlisting}
%b = [[1, 2, 3]]
%np.repeat(b, 3, axis = 0)
%array([[1, 2, 3],
%[1, 2, 3],
%[1, 2, 3]])
%\end{lstlisting}
%In general, be careful about having the correct rows, columns, transposes, etc throughout your code.

Now we maximize over choices of $W_{T+1}$ (choosing how much to save for the next period).  Then we will have a row vector
representing the value function for period $T$ across all possible $W_{T+1}$.  Iterate this procedure to fill in the value
function for all $t=T+1,T,\ldots, 0$.

\item In each iteration, you maximize to find the value function at time $t$.  Save the values of $W_{t+1}$ that achieve the
maximum.  The result is an $N$ by $T+1$ matrix whose $(n,t)$ entry gives the optimal amount of cake to leave for period
$t+1$, given that we start period $t$ with the the $n$-th value of our vector of cake.  This is the policy function.

\item If the keyword argument \li{plot} is set to \li{True}, plot the surface of the Value and Policy functions.
This can be done by including the following import lines
\begin{lstlisting}
>>> from matplotlib import pyplot as plt
>>> from matplotlib import cm
>>> from mpl_toolkits.mplot3d import Axes3D
\end{lstlisting}
and using the following code:
\begin{lstlisting}
>>> W = np.linspace(0, Wmax, N)
>>> x = np.arange(0, N)
>>> y = np.arange(0, T+2)
>>> X, Y = np.meshgrid(x, y)
>>> fig1 = plt.figure()
>>> ax1 = Axes3D(fig1)
>>> ax1.plot_surface(W[X], Y, np.transpose(V), cmap=cm.coolwarm)
>>> plt.show()

>>> fig2 = plt.figure()
>>> ax2 = Axes3D(fig2)
>>> y = np.arange(0,T+1)
>>> X, Y = np.meshgrid(x, y)
>>> ax2.plot_surface(W[X], Y, np.transpose(psi), cmap=cm.coolwarm)
>>> plt.show()
\end{lstlisting}
where \li{W} is the vector of cake amounts, \li{V} is the value function, and \li{psi} is the policy function.

\item Return the arrays giving the value function and the policy function.
\end{enumerate}

Solve the problem using cake size $1$, discount factor $\beta = .9$, number of time periods $T = 10$, and number of
discrete cake values $N = 100$. You should also try plotting the value and policy functions for fixed time periods across
$W_t$, or for fixed $W_t$ across time,
 and make sure that these plots fit your intuition. See Figure \ref{fig:valueslices}. Your output should
agree with the figure.
\end{problem}

\begin{figure}
\begin{subfigure}{.5\textwidth}
    \includegraphics[width=\textwidth]{fixed_time.pdf}
\end{subfigure}
\begin{subfigure}{.5\textwidth}
    \includegraphics[width=\textwidth]{fixed_w.pdf}
\end{subfigure}
\caption{Slices of the finite horizon value function for fixed values of $t$ and $W$, respectively.}
\label{fig:valueslices}
\end{figure}

\section*{The Recursive Problem, Infinite Horizon}\label{SecRecProbInFHor}
Next we consider an infinite horizon problem.  For simplicity, we continue with the example from the previous section.
Suppose that rather than optimizing over $t = 0,1,\ldots,T$, we wish to optimize over an infinite time horizon:
\begin{equation*}
\text{maximize} \sum_{t=0}^\infty \beta^t u(W_t,c_t) \quad \text{s.t.} \quad W_{t+1} = g(W_t,c_t).
\end{equation*}
Since at any time $t$, there are an infinite number of periods remaining, one might suspect that the optimal policy will
not depend on the current time $t$.

\begin{problem}
\label{prob:cake_prob2}
Compute the solution to Problem \ref{prob:cake_prob} with $T = 1000$, and the rest of the inputs the same.
Plot the policy function across time for fixed $W_t = 1$.
Notice that it is the same for all time periods, except those near the end time $T$.
\end{problem}

As suggested by the results of Problem \ref{prob:cake_prob2},  the policy function for the infinite horizon problem does not
depend on the time $t$ (this can be proved).  That is, at any time $t$, the optimal decision depends only on the amount of cake
at the beginning of the period, not the value of $t$.  So everything can now be written in terms of variables today and variables
tomorrow. We will denote variables tomorrow with a ``$\:'\:$".
\begin{equation}
\label{EqBellman}
V\left(W\right) = \max_{W'\in[0,W]}\:\: \left\{u\left(W - W'\right) + \beta V\left(W'\right)\right\}
\end{equation}
Note that the value function $V$ on the left-hand-side of \eqref{EqBellman} and on the right-hand-side are the same function.

Because the problem now has an infinite horizon, the nature of the solution is a little different. The solution to \eqref{EqBellman}
is a policy function $W'=\psi(W)$ that creates a fixed point in $V$. In other words, the solution is a policy function $\psi(W)$
that makes the function $V$ on the left-hand-side of \eqref{EqBellman} equal the function $V$ on the right-hand-side.

Define $C$ as an operator on any value function $V_k\left(W\right)$. Let $C$ perform the following operation.
\begin{equation}
\label{EqContraction}
C\Bigl(V_k\left(W\right)\Bigr) \equiv \max_{W'\in[0,W]}\:\: \left\{u\left(W-W'\right) + \beta V_k\left(W'\right)\right\}.
\end{equation}
Note that the value function on the right-hand-side of \eqref{EqContraction} and on the left-hand-side are the same function $V_k$,
but have different inputs--$W$ versus $W'$. The operator $C$ takes in a function $V_k$, and gives a new
function which we will call $V_{k+1}$:
\begin{equation*}
V_{k+1}\left(W\right) \equiv C\Bigl(V_k\left(W\right)\Bigr).
\end{equation*}
The value function $V_{k+1}$ that results from the operation $C$ is not necessarily the same as the value function that the system
began with ($V_k$). However, according to equation \eqref{EqBellman} we seek a $V$ such that $C(V) = V$.  The solution, then, is the fixed point in $V$.
\begin{equation*}
C\Bigl(V_k\left(W\right)\Bigr) = V_{k+1}\left(W\right) = V_k\left(W\right) = V\left(W\right)
\end{equation*}

When trying to solve a fixed point equation, it is often very helpful to utilize the Contraction Mapping Principle, which
guarantees the existence of a fixed point of a mapping, provided that the map sends any two distinct inputs to outputs
that are strictly closer to each other than the inputs, in a controlled way. This principle also provides a constructive
way to obtain the fixed point, namely by iterating the map.
Fortunately, it can be shown that if $u(\cdot)$ is real-valued, continuous, and bounded, $\beta\in(0,1)$, and that the constraint
set $W'\in[0,W]$ is nonempty, compact-valued, and continuous, then the operator $C$ is a contraction and thus we can obtain
a solution $V$ by iteration:
\begin{equation*}
\lim_{k\rightarrow\infty}\: C^k\Bigl(V_0\left(W\right)\Bigr) = V_k(W) =  V\left(W\right)
\end{equation*}
for any $V_0$.

Remember, in the infinite horizon problem both the value and policy functions do not depend on time.  Computationally, this means that
 the value and policy functions in the infinite horizon problem are one dimensional.

\begin{problem}
Expand your \li{eatCake} function to solve the Cake Eating Problem with an infinite time horizon. If the keyword argument
\li{finite} has the value \li{True}, then your function should behave as in Problem \ref{prob:cake_prob}, solving the finite
time horizon problem. However, if \li{finite = False}, solve the infinite time horizon problem through the following steps.
Both problems will require you to pre-compute the values $u(W - W')$, where $W$ and $W'$ range over the set of discrete cake
amounts. Be sure to avoid replicating code by factoring it out.
As in Problem \ref{prob:cake_prob}, take $u(c_t) = \sqrt{c_t}$.
\begin{enumerate}
\item As in Problem \ref{prob:cake_prob}, approximate the continuum of possible
cake sizes by a column vector called $W$ that ranges from 0 to $W_{max}$ in $N$ steps.

\item \label{item:step2} Initialize the value function V as a vector of zeros of length $N$.  This is $V_0$.  Perform one iteration
of the contraction operation given in equation \eqref{EqContraction} to get a new value function $V_1$ (this should be very similar
to Problem 1).  Determine the resulting policy function $W' = \psi_1\left(W\right)$.  [HINT: The policy function should be a vector
of length $N$ of optimal future values of the cake $W'$ given the current value of the cake $W$, and $V_T$ should be an $N$-length
vector representing the value of entering a period with cake size $W$.]

\item \label{item:step3} Measure the distance between the two value functions as the sum of the
squared differences,
\begin{equation}
\label{EqDist}
\delta_1\equiv \norm{V_1\left(W\right) - V_0\left(W'\right)}_2^2 = \left(V_1 - V_0\right)^T\left(V_1 - V_0\right).
\end{equation}
Defined in this way, $\delta_1\in [0,\infty)$.

%\item \label{item:step4} Take the resulting $V_1$ from \ref{item:step2}, and perform the same contraction on it to generate $V_2$
%and $\psi_2$. That is, generate,
%\begin{equation*}
%  V_2\left(W\right) = C\Bigl(V_1\left(W\right)\Bigr) = \max_{W'\in[0,W]}\: u\left(W - W'\right) + \beta V_1\left(W'\right)
%\end{equation*}
%and the accompanying policy function $W'=\psi_2\left(W\right)$. Calculate the accompanying distance measure for $\delta_2$ using
% the formula from \eqref{EqDist} with the updated period subscripts. Compare $\delta_2$ with $\delta_1$ from \ref{item:step3}.
%
%\item \label{item:step5} Repeat \ref{item:step4} and generate $V_3$ and $\psi_2$ by performing the contraction on $V_2$. Compare
%$\delta_3$ to $\delta_2$ and $\delta_1$.

\item Write a loop that performs the contraction operation from steps \ref{item:step2} and \ref{item:step3} iteratively
until the distance measure is very small ($\delta_k < 10^{-9}$).  The distance measure $\delta_k$ being arbitrarily close to zero means
 you have converged to the fixed point $V_k = V_{k+1} = V$. (For fun, you can show that the policy function converges to the same
 function regardless of what you put in for your initial policy function value.)

\item If \li{plot = True}, plot the converged policy function vector
($y$-axis) as a function of the cake amounts ($x$-axis).

\item Return the value function and policy function arrays.

Compute the value function and policy function for the infinite time horizon problem with
cake size $1$, discount factor $\beta = .9$, and number of
discrete cake values $N = 100$. The plot you generate should agree with Figure \ref{fig:infinitePolicy}.
\end{enumerate}
\end{problem}

\begin{figure}
\includegraphics[width=\textwidth]{infiniteHorizon.pdf}
\caption{Policy function for infinite time horizon.}
\label{fig:infinitePolicy}
\end{figure}

\begin{figure}
\includegraphics[width=\textwidth]{convergence.pdf}
\caption{Due to the contraction mapping principle, $\delta_k$ decreases as we perform iterations
until it is small enough to meet our convergence tolerance.}
\end{figure}

\section*{Infinite Horizon, Stochastic, i.i.d.}\label{SecRecProbInfinHorStochiid}

In practice, dynamic programming problems often involve some level of uncertainty.
For example, as time progresses prices may fluctuate, resources may vary, or preferences themselves may change.
 In this lab, we reexamine the cake eating problem, this time allowing for uncertainty.

We consider again the problem of optimizing a sequence of decisions over an infinite time horizon.
We assume that the individual's preferences deviate each period according to some ``shock" $\ve$,
where $\ve$ is a random variable.  We assume that the shock terms $\ve$ for each time period are
independent and identically
distributed (i.i.d.).  In effect, this means the probabilities associated with the $\ve$ are the
same for any time $t$ and do not depend on each other.  We assume for now that the $\ve$ are
distributed normally with mean $\mu$ and variance $\sigma^2$.  The Bellman equation can be easily
rewritten in the following way to incorporate the uncertainty,
\begin{equation}\label{stoch_Bellman}
   V\left(W,\ve\right) = \max_{W'\in[0,W]}\: \{\ve u\left(W - W'\right) +
   \beta E_{\ve'}\left[V\left(W',\ve'\right)\right]\},
\end{equation}
where $\ve \sim N(\mu,\sigma^2)$ and
$E$ is the unconditional expectation operator over $\ve$.  Note that now the value function depends
on two variables.  It represents the value of entering the period with $W$, the amount of cake,
and a preference shock of $\ve$.  For example, in a period where the realization of $\ve$ is higher,
we will get more value from the cake eaten in the current period.  Because we do not know the value of
the shock in the next period $\ve'$, we consider only the expected value for future time.

As it turns out, we can solve this problem in a manner similar to the infinite horizon deterministic cake-eating
problem considered in the Value Function Iteration lab.  It is worth noting that in this case,
 the value and policy functions will be two dimensional, as they will depend on both $W$ and $\ve$.

In order to deal with $\ve$ computationally, we would like to represent it as a vector of possible values
it could take, along with the corresponding probabilities that it takes each of those values.  However,
$N(\mu,\sigma^2)$ is a continuous distribution, so we cannot represent every value $\ve$ could take.
We need a discrete distribution that approximates $N(\mu,\sigma^2)$.

To do this, we choose $N$ equally spaced points centered about the mean at which to approximate the
distribution. Call these values $\ve_1,\ldots,\ve_N$, and let the spacing between adjacent points be
given by $\delta$.
We can then break up the support of the distribution into $N$ bins, where adjacent bins share a common
endpoint. Call the endpoints of these bins $v_1,\ldots,v_{N+1}$. By choosing the endpoints to be
halfway between each $\ve_k$, we have the formula
\[
v_k = \ve_k - \frac{1}{2}\delta, \qquad k=1,\ldots,N
\]
and
\[
v_{N+1} = \ve_N + \frac{1}{2}\delta.
\]

\begin{figure}[h!]
\label{stoch1_fig1}
\begin{center}
\includegraphics[width = \textwidth]{discnorm.pdf}
\end{center}
\caption{Discretization of $N(\mu,\sigma^2)$.  We approximate $P(\ve = \ve_k)$ by the area of the shaded region.}
\end{figure}

We can then associate $\ve_k$ with the area under the curve from $v_k$ to $v_{k+1}$.
In Python, we can find the area using the function \li{norm.cdf} found in the \li{stats} package.
The cdf (cumulative distribution function) gives the area under the curve from $-\infty$ to a specified value.
For example, in the following code, \li{eps} is the area under the curve from 0 to 1.

\begin{lstlisting}
>>> from scipy import stats as st
>>> mu = 0
>>> sigma = 1
>>> eps = st.norm.cdf(1,loc=mu,scale=sigma) - st.norm.cdf(0,loc=mu,scale=sigma)
\end{lstlisting}

In general, it is sufficient to take our points $\ve_k$ ranging from $\mu - 3\sigma$ to $\mu + 3\sigma$, as this
range contains about 99.7\% of the probability mass.

\begin{problem}
Write a function called \li{discretenorm} that accepts an integer $K$ representing the number of discrete points
desired, a mean $\mu$, and a standard deviation $\sigma$. It should return a length-$K$ vector of equally-spaced
values ranging from $\mu - 3\sigma$ to $\mu + 3\sigma$ inclusive,
and a length-$K$ vector containing the associated probabilities.
Plot the approximation of $N(0,1)$ using different values of $K$ to check that your results are plausible.
\end{problem}

Now that we have a discrete distribution for $\ve$, we can solve for the value and policy functions
determined by \eqref{stoch_Bellman}.

\begin{problem}
Complete the following steps to solve the problem described above.
Assume that the period utility function is $u(c)=\sqrt{c}$.
Write a function \li{stochEatCake}
that accepts parameters $\beta$ (discount factor), $N$ (number of discrete cake values),
a tuple of values \li{e_params}, $W_{max}$ (the original size of the cake, set to default value 1),
a keyword argument \li{iid} (set to default value \li{True}), and
a keyword argument \li{plot} (set to default value of \li{False}). Inside the function, carry out the steps
outlined below.

The argument \li{e_params} is a tuple consisting of the values needed to generate
the discrete approximation to $\ve$. In the present case, this tuple consists (in order) of
$K$ (the number of discrete approximations of $\ve$), $\mu$ (the
mean of the shock term $\ve$), and $\sigma$ (the standard deviation of the shock term $\ve$),
since these are the arguments we need to pass to our \li{discretenorm} function.

\begin{enumerate}
\item First, compute an approximation of $\ve$ using the \li{discretenorm} function created in Problem 1.
Use $K$ equally spaced points to approximate $N(\mu,\sigma^2)$. Denote the resulting $K$-length
vector of equally-spaced values by
\[e =(e_1,\ldots,e_K),
\]
and denote the $K$-length vector of the associated probabilities
by
\[\Gamma = (\Gamma_1,\ldots,\Gamma_K).
\]
Note that $\Gamma_k$ give the probability $P(\ve = e_k)$.

Since the values needed for the \li{discretenorm} function are contained in the \li{e_params} input,
we can feed these values directly into the function in the following way:
\begin{lstlisting}
>>> e, gamma = discretenorm(*e_params)
\end{lstlisting}
The \li{*} operator essentially unpacks the values of a tuple or list.

\item As done previously, create a vector
\[w = (w_1,\ldots,w_N)
\]
of possible cake sizes. This should be
a length-$N$ vector of equally spaced values from 0 to $W_{max}$, inclusive.

\item Represent the value function as a $N \times K$ matrix $v$, satisfying
\[
v_{i,j} = V(w_i, e_j).
\]
(The rows correspond to different values of $W$ and the columns correspond to different values of $\ve$.)
Initialize each entry of the matrix to 0.

Likewise, represent the policy function as a $N \times K$ matrix $p$, satisfying
\[
p_{i,j} = \psi(w_i,e_j).
\]
Initialize all entries to 0.

\item In order to evaluate the value function equation, we need to pre-compute $\ve u(W-W')$ for all values of
$\ve,W,W'$.
Begin by computing all possible values of $u(W-W')$, and storing these values in a $N \times N$ array,
as done before. Call this array $u$. Make sure that the upper triangular
entries of this array are equal to zero, as these entries correspond to consuming more cake than is
available, which is impossible.

The values $\ve u(W-W')$ will be represented by a three-dimensional array $\hat{u}$ of size
$N\times N\times K$, satisfying
\[
\hat{u}_{i,j,k} = v_{i,j}e_k.
\]
We can compute this array easily as follows:

\begin{lstlisting}
>>> import numpy as np
>>> u_hat = np.repeat(u, K).reshape((N,N,K))*e
\end{lstlisting}

\item We also need to compute $E_{\ve'}\Bigl[V\left(W',\ve'\right)\Bigr]$ for each value of $W'$.
The expected value is simply
\begin{equation*}
E_{\ve'}\Bigl[V\left(W',\ve'\right)\Bigr] = \sum_{k=1}^K \Gamma_kV(W',e_k').
\end{equation*}
The result is a length $N$ vector, call it $E$, satisfying
\[
E_i = E_{\ve'}\Bigl[V\left(w_i,\ve'\right)\Bigr] = \sum_{k=1}^K \Gamma_kv_{i,k}
\]
This calculation can be done by multiplying $\Gamma$ element-wise to each row of the
value function matrix $v$, and then summing along the rows. Something like the following
line of code should do the trick:
\begin{lstlisting}
>>> E = (v*gamma).sum(axis=1)
\end{lstlisting}

\item We can now compute the value function contraction
\begin{equation*}\label{EqContractStochiid}
C\Bigl(V\left(W,\ve\right)\Bigr) \equiv \max_{W'\in[0,W]}\:
\Bigl\{\ve u\left(W-W'\right) + \beta E_{\ve'}\Bigl[V\left(W',\ve'\right)\Bigr]\Bigr\}.
\end{equation*}
The first task is to create an $N \times N \times K$ array $c$ satisfying
\[
c_{i,j,k} = \hat{u}_{i,j,k} + \beta E_j.
\]
This can be done in any manner of ways. Below is a one-liner that does the job.
\begin{lstlisting}
>>> c = np.swapaxes(np.swapaxes(u_hat, 1, 2) + beta*E, 1, 2)
\end{lstlisting}

Now, for any $k$, for all $i < j$, set $c_{i,j,k}$ to a large negative number, say $-10^{10}$,
so that when maximizing over this array, we do not choose to consume more cake than is available.
Again, this can be done in a variety of different ways, but the following does the job concisely:
\begin{lstlisting}
>>> c[np.triu_indices(N, k=1)] = -1e10
\end{lstlisting}

Finally, maximize over the second axis of $c$ (which corresponds to different values of $W'$)
to obtain the updated value function matrix:
\begin{lstlisting}
>>> v_new = np.max(c, axis=1)
\end{lstlisting}
You can likewise update your policy function matrix as follows:
\begin{lstlisting}
>>> max_indices = np.argmax(c, axis=1)
>>> p = w[max_indices]
\end{lstlisting}

\item We now have our updated value function matrix $v_{new}$ as well as the
previous $v$, which we refer to here as $v_{old}$. As we iterate on the value function equation, we need a norm
\begin{equation*}
\delta = \|v_{new} - v_{old}\|_2
\end{equation*}
that measures the distance between these two value functions to determine convergence.
You may compute the norm using the SciPy function \li{scipy.linalg.norm}, or by direct calculation.
At the end of each iteration, make sure to set $v$ to $v_{new}$, so that the updates carry through the
loop.
Iterate on the contraction until $\delta < 10^{-9}$.

\item If \li{plot = True}, make a 3-D surface plot of the policy function for the converged problem
$W' = \psi\left(W,\ve\right)$ which gives the value of the cake tomorrow as a
function of the cake today  and the taste shock today.  Do the same for the value function.
Example code to create the value function plot is provided below.
\begin{lstlisting}
>>> x = np.arange(0,N)
>>> y = np.arange(0,K)
>>> X,Y = np.meshgrid(x,y)
>>> fig1 = plt.figure()
>>> ax1 = Axes3D(fig1)
>>> ax1.plot_surface(w[X], Y, v.T, cmap=cm.coolwarm)
>>> plt.show()
\end{lstlisting}
Creating the policy function plot is similar.

\item Return the converged value function matrix $v$ and policy function matrix $p$.

\end{enumerate}
Test your function using values $\beta = .9$, $N = 100$, $K = 7$, $\sigma = .5$, $\mu = 4\sigma$,
and \li{plot = True}.
The proper way to set this up and call the function is as follows:
\begin{lstlisting}
>>> e_params = (7, 4*.5, .5)
>>> stuff = stochEatCake(.9, 100, e_params, plot=True)
\end{lstlisting}
\end{problem}

\begin{figure}
    \centering
    \includegraphics[width = \textwidth]{stoch_value.pdf}
    \caption{3D surface representing the value function for the Stochastic Cake-Eating problem.}
\end{figure}

\section*{Infinite Horizon, Stochastic, AR(1)}\label{SecRecProbInfinHorStochAR1}

In the previous example, we assumed that the shocks at time $t$ were independent of what happened in
previous periods.  Often a shock may depend on recent events.  We will assume now that the shocks are persistent,
meaning preferences in the current period are more likely to be close to what they were in the previous period.
We can characterize the persistence by what is called an autoregressive process of order one, denoted AR(1).
Such a process is defined as follows.
\begin{equation}\label{EqAR1shock}
\ve' = (1-\rho)\mu + \rho\ve + \nu' \quad\text{where}\quad \rho\in(0,1) \quad\text{and}\quad \nu\sim N(0,\sigma^2).
\end{equation}

Essentially, instead of allowing the shocks to have a mean which is independent of the past, the mean is now a
weighted average (weighted by $\rho$) of some $\mu$ and the previous realization of the shock, $\ve$.  As it turns
 out, we can approximate this process by thinking of it as a Markov Chain. This means we need to determine a
 discrete set of points representing possible values of $\ve$ and a Markov transition matrix that gives the
 probabilities of moving from one value of $\ve$ to another.  There are methods for determining the discrete
 approximation of $\ve$ with a Markov transition matrix.  These methods are beyond the scope of this section,
 but you can use the file \li{tauchenhussey.py} to implement them in the next problem.

The Bellman equation becomes the following, in which the only change from the i.i.d. shock case is that the
expectations operator is now conditional on the current shock $\ve$:

\begin{equation*}
   V\left(W,\ve\right) = \max_{W'\in[0,W]}\: \{\ve u\left(W - W'\right) +
   \beta E_{\ve'|\ve}\left[V\left(W',\ve'\right)\right]\},
\end{equation*}
where $\ve'$ is distributed according to \eqref{EqAR1shock}.
Let $\Gamma_{i,j}=P\left(\ve_j'|\ve_i\right)$ where $\ve_j'$ is the value of the shock in
the next period and $\ve_i$ is the value of the shock in the current period.
In other words, $\Gamma$ is the Markov transition matrix.

The solution to this problem is of the same type as that in the i.i.d. case, since the only difference is
the probability distributions of the $\ve$.

\begin{problem}
Expand your \li{stochEatCake} function to handle the case of AR(1) shock terms. The function should
handle this case for the parameter value \li{iid = False}, and should handle the previous case of
normally distributed i.i.d. shock terms for the parameter value \li{iid = True}. You will need to
add a few ``if ... else" statements, as well as implement the steps outlined below, but most of the
code will remain unchanged.

\begin{enumerate}
\item In the AR(1) case, the \li{e_params} argument should be a tuple of values needed to
generate the arrays $e$ and $\Gamma$ that approximate the values and distribution of $\ve$
as a Markov chain.
Use the file \li{tauchenhussey.py} to calculate these arrays.
The provided Python function \li{tauchenhussey} produces the vector $e$ of length $M$
and an $M\times M$ transition matrix $\Gamma$.
Thus, you simply need the following lines of code, similar to the previous case.
\begin{lstlisting}
>>> from tauchenhussey import tauchenhussey
>>> e, gamma = tauchenhussey(*e_params)
\end{lstlisting}

\item Because our values for $e$ and $\Gamma$ are different in the AR(1) case than
in the i.i.d. case, we must compute the expectation in a different manner.
In particular, we need to compute the conditional expectation
\begin{equation*}
E_{\ve'|\ve}\Bigl[V\left(W',\ve'\right)\Bigr].
\end{equation*}
We obtain a two-dimensional array, since the expectation depends on both $W'$ and on $\ve$.
The expectation can be computed by the matrix multiplication $v\Gamma^T$.
Your code should match the following.
\begin{lstlisting}
>>> E = v.dot(gamma.T)
\end{lstlisting}

\item The last difference comes in computing the array $c$. Fortunately, it is easier in this case.
Recall that $c$ gives the values for
\[
\ve u\left(W-W'\right) + \beta E_{\ve'|\ve}\Bigl[V\left(W',\ve'\right)\Bigr].
\]
The array $\hat{u}$ contains the values for the first term in the expression, and the array $E$
contains the values for the expectation term.
Hence, we obtain $c$ by simple addition. Array broadcasting makes this work without problems.
\begin{lstlisting}
>>> c = u_hat + beta*E
\end{lstlisting}
You will still need to set the upper triangular entries of $c$ to a large negative number, just as in the
previous case.
\end{enumerate}

Those are the only differences. Let the following code snippet be a guideline for how to implement
these differences.
\begin{lstlisting}
>>> if iid:
>>>     # compute E as outlined in the previous problem
>>> else:
>>>     # compute E as outlined in the current problem
\end{lstlisting}

Now test your function with $\beta = .9$, $N = 100$, \li{iid = False}, and \li{plot = True}.
As inputs to \li{tauchenhussey}, let $K=7$, the mean of the process
$\mu=4\sigma$, $\rho = 1/2$, $\sigma=1/2$, and
\[baseSigma=(0.5+\frac{\rho}{4})\sigma +
(0.5 - \frac{\rho}{4})\frac{\sigma}{\sqrt{1-\rho^2}}.
\]
Your \li{e_params} parameter will therefore be a tuple of values containing (in order)
$K$, $\mu$, $\rho$, $\sigma$, and $baseSigma$.
\end{problem}
\end{comment}