\documentclass{article}[11pt]
\usepackage{amsmath}
\usepackage{amsfonts}
\usepackage{enumitem}
\usepackage{mathtools}
\usepackage{mathrsfs}
\usepackage{amssymb}
\usepackage{amsthm}
\usepackage{graphicx}
\usepackage{adjustbox}
\usepackage{fixmath}
\usepackage{xcolor}
\usepackage{float}
\usepackage{caption}
\newcommand*{\Z}{\mathbb{Z}}
\newcommand*{\D}{\mathbb{D}}
\newcommand*{\Q}{\mathbb{Q}}
\newcommand*{\R}{\mathbb{R}}
\newcommand*{\N}{\mathbb{N}}
\newcommand*{\I}{\mathbb{I}}
\newcommand*{\C}{\mathbb{C}}
\newcommand*{\Prime}{\mathbb{P}}
\newcommand*{\E}{\textrm{E}}
\newcommand*{\Var}{\textrm{Var}}
\newcommand*{\m}{\textrm{mod }}
\newcommand*{\osc}{\textrm{osc}}
\newcommand*{\vol}{\textrm{vol}}
\newcommand*{\prt}{\mathcal{P}}
\newcommand*{\intr}{\textrm{int}}
\newcommand*{\cl}{\textrm{cl}}
\newcommand*{\bd}{\textrm{bd}}
\newcommand*{\tr}{\textrm{Tr}}
\newcommand*{\img}{\textrm{im}}
\newcommand*{\rank}{\textrm{rank}}
\newcommand*{\spn}{\textrm{span}}
\newcommand*{\Trig}{\textrm{Trig}}
\newcommand*{\col}{\textrm{col}}
\newcommand*{\iu}{\mathbold{i}}
\newcommand*{\re}{\textrm{Re}}
\newcommand*{\im}{\textrm{Im}}
\newcommand*{\e}{\textrm{e}}
\newcommand*{\arcsinh}{\textrm{arcsinh}}
\newcommand*{\ros}{\widehat{\mathcal{R}}}
\newcommand*{\defeq}{\vcentcolon=}
\newcommand*{\eqdef}{=\vcentcolon}
\DeclareMathOperator*{\diam}{diam}
\DeclareMathOperator*{\Bor}{Bor}
\linespread{1.3}

\renewcommand{\dot}{\;\begin{picture}(-1,1)(-1,-3)\circle*{3}\end{picture}\;\,\, }
\newcommand\norm[1]{\left\lVert#1\right\rVert}
\newcommand*{\QED}{\hfill\ensuremath{\square}}%

\newtheoremstyle{dotless}{}{}{}{}{\bfseries}{}{12pt}{}

\theoremstyle{dotless}
\newtheorem{thm}{Theorem}[section]
\newtheorem{lem}[thm]{Lemma}
\newtheorem{cor}[thm]{Corollary}
\newtheorem{prop}[thm]{Proposition}
\newtheorem{defn}[thm]{Definition}
\newtheorem{rem}[thm]{Remark}
\newtheorem{exa}[thm]{Example}
\newtheorem{exe}[thm]{Exercise}
\newtheorem{aside}[thm]{Aside}

% stuff from the percent sign to end of line is a comment, ignored by LaTeX

\usepackage[margin=1in]{geometry} %set margins
\begin{document}

\nocite{*} % this command forces all references in template.bib to be printed in the bibliography

\title{Hausdorff Measure and Dimension}

\author{Name \\
Course Number – Course Name \\
University Name}

\maketitle

\section*{Introduction}

The \textbf{Hausdorff measure} is a measure on metric spaces which generalizes the Lebesgue measure on $\R^n$.
As such, one can imagine how useful this measure can be in areas such as differential geometry, where it
can be used to define ``volumes'' of submanifolds (such as the unit sphere in $\R^3$).
In \textbf{Assignment 3}, we've seen how to construct the Hausdorff measure on a metric space. We shall now look at
some further properties of the measure (such as its \textit{measurable} sets) as well as introduce
\textbf{Hausdorff dimension}, which generalizes the notion of the usual (topological) dimension.

\section{Measurable sets of the Hausdorff measure}

Let $(X, \rho)$ be a metric space.
Recall (from \textbf{Assignment 3}) that if $d > 0$ and $\delta > 0$, we define $H_\delta^d : \mathcal{P}(X) \to [0,\infty]$ by:
\[ H_\delta^d(E) = \inf \left\{ \sum_{i \geqslant 1} \diam(U_i)^d : E \subseteq \bigcup_{i \geqslant 1} U_i, \; \diam(U_i) < \delta \right\} \]

\noindent
for all $E \subseteq X$,
where $\diam(U) = \sup \{ \rho(x,y) : x,y \in U \}$. Then the $d$-dimensional Hausdorff \textit{outer} measure $H^d$ on $(X, \rho)$ is defined by:
\[ H^d(E) = \lim_{\delta \to 0^+} H_\delta^d(E) \]

\noindent
for all $E \subseteq X$. That is, $H^n(E)$ is approximated by
the sum of $d$\textsuperscript{th} powers of the diameters of the sets of countable covers of $E$ where each diameter is less than $\delta$,
and a smaller value of $\delta$ means a better
approximation. We wish to find the \textit{measurable} sets of $H^d$.
\bigskip

We know that a set $E \subseteq X$ is $H^d$-measurable when $H^d(A) = H^d(E \cap A) + H^d(A \setminus E)$
for all $A \subseteq E$. We therefore have:
\[ \mathcal{B} = \{ E \subseteq X : H^d(A) = H^d(E \cap A) + H^d(A \setminus E) \; \forall A \subseteq X \} \]

\noindent
is the set of all $H^d$-measurable sets. In particular, the Carath\'{e}odory theorem tells us that $\mathcal{B}$ is a $\sigma$-algebra
and that $H^d$ restricted to $\mathcal{B}$ is a \textit{complete} measure. Of course, this definition doesn't tell us much about exactly
which sets are $H^d$-measurable. If we wish for the Hausdorff measure to extend the notion of the Lebesgue measure, it's reasonable
to ask that
the Borel sets in $X$ be $H^d$-measurable. Recall (again from \textbf{Assignment 3}) that two sets $A, B \subseteq X$ in
a metric space $(X,\rho)$ are said to be \textbf{positively separated} when $\inf \{ \rho(x,y) : x \in A, \; y \in B \} > 0$.
For brevity, we will define $\rho(A,B) = \inf \{ \rho(x,y) : x \in A, \; y \in B \}$. We have seen that if $A$ and $B$ are
positively separated, then $H^d(A \cup B) = H^d(A) + H^d(B)$.

\begin{defn}
	In general, if $\mu^*$ is an outer measure on a metric space $(X, \rho)$ such that $\mu^*(A \cup B) = \mu^*(A) + \mu^*(B)$
	whenever $A, B \subseteq X$ are positively separated, then $\mu^*$ is called a \textbf{metric outer measure} (or a
	\textbf{Carath\'{e}odory outer measure}).
\end{defn}
\bigskip

Hence, the Hausdorff outer measure $H^d$ is a metric outer measure. We have the following general result about metric outer
measures that states that the Borel sets must be measurable.

\begin{prop}
	Let $(X, \rho)$ be a metric space and $\mu^*$ a metric outer measure on $(X, \rho)$. Then
	every Borel set in $X$ is $\mu^*$-measurable.
\end{prop}
\begin{proof}
	(Adapted from \cite{Fol}).
	\bigskip

	\noindent
	Let $\Bor(X)$ denote the set of Borel sets, which we know is a $\sigma$-algebra. In particular, we have that
	$\Bor(X) = \sigma(\{ F \subseteq X : F \text{ is closed} \})$. Let $\mathcal{B}$ be the $\sigma$-algebra of $\mu^*$-measurable
	sets in $X$. Then to show that $\Bor(X) \subseteq \mathcal{B}$, it suffices to prove that every closed set is $\mu^*$-measurable.
	\bigskip

	\noindent
	Indeed, let $F \subseteq X$ be closed and let $A \subseteq X$ be any set. It's immediate that
	$\mu^*(A) \leqslant \mu^*(A \cap F) + \mu^*(A \setminus F)$ by subadditivity. Moreover, if
	$\mu^*(A) = \infty$, equality clearly holds. Thus, we may assume $\mu^*(A) < \infty$.
	For each $n \in \N$, let $A_n = \{ x \in A \setminus F : \rho(\{x\}, F) = \inf\{ \rho(x,y) : y \in F \} \geqslant \frac{1}{n} \}$.
	Observe that if $x \in A_n$, then $\rho(\{x\}, F) \geqslant \frac{1}{n} > \frac{1}{n+1}$, so $x \in A_{n+1}$.
	Hence, $A_1 \subseteq A_2 \subseteq \cdots \subseteq A \setminus F$.
	\bigskip

	\noindent
	Now for any $x \in A \setminus F$, suppose we had that $\rho(\{x\}, F) = \inf\{ \rho(x,y) : y \in F \} = 0$.
	Then for all $n \in \N$, we can find some $y_n \in F$ for which $\rho(x,y_n) < \frac{1}{n}$. As $\rho$ is a metric
	on $X$, it follows that $y_n$ converges to $x$. By closure of $F$, we have that in fact $x \in F$, which is a contradiction.
	It thus follows that $\rho(\{x\}, F) > 0$ for any $x \in A \setminus F$, which means we can find a $k \in \N$ for which
	$\rho(\{x\}, F) \geqslant \frac{1}{k}$ and consequently $x \in A_k \subseteq \bigcup_{n \geqslant 1} A_n$. This proves
	that $\bigcup_{n \geqslant 1} A_n = A \setminus F$.
	\bigskip

	\noindent
	As $A = (A \cap F) \cup (A \setminus F) \supseteq (A \cap F) \cup A_n$ for all $n \in \N$, we have by monotonicity that:
	\[ \mu^*(A) \geqslant \mu^*((A \cap F) \cup A_n) \]

	\noindent
	Moreover, we have by construction that $\rho(\{x\}, F) \geqslant \frac{1}{n}$ whenever $x \in A_n$, and thus
	(in particular) $\rho(A_n, A \cap F) \geqslant \frac{1}{n} > 0$, so $A_n$ and $A \cap F$ are positively separated.
	Since $\mu^*$ is a \textit{metric} outer measure:
	\[ \mu^*((A \cap F) \cup A_n) = \mu^*(A \cap F) + \mu^*(A_n) \]

	\noindent
	Now consider $D_n = A_{n+1} \setminus A_n = \{ x \in A \setminus F : \frac{1}{n+1} \leqslant \rho(\{x\}, F) < \frac{1}{n} \}$.
	By construction, we have that the $D_n$'s are disjoint and $A \setminus F = \bigcup_{n \geqslant 1} D_n$.
	Let $x \in D_{n+1}$ and suppose $y \in X$ satisfies $\rho(x,y) < \frac{1}{n(n+1)}$. Then for all $z \in F$, $x \in D_{n+1} = A_{n+2} \setminus
	A_{n+1}$ means $\frac{1}{n+2} \leqslant \rho(x,z) < \frac{1}{n+1}$, so the triangle inequality gives:
	\[ \rho(y,z) \leqslant \rho(y,x) + \rho(x,z) < \frac{1}{n(n+1)} + \frac{1}{n+1} = \frac{n+1}{n(n+1)} = \frac{1}{n} \]

	\noindent
	and thus $\rho(y,F) < \frac{1}{n} \implies y \not\in A_n$. This means if $y \in A_n$, then we must have
	$\rho(x,y) \geqslant \frac{1}{n(n+1)}$ so that $\rho(x, A_n) \geqslant \frac{1}{n(n+1)}$. As our choice of $x \in D_{n+1}$ was arbitrary,
	it follows that $\rho(D_{n+1}, A_n) \geqslant \frac{1}{n(n+1)} > 0$ so that $D_{n+1}$ and $A_n$ are positively separated.
	Thus, for all $n \in \N$:
	\[ \mu^*(A_{2n+1}) = \mu^*(D_{2n} \cup A_{2n}) \geqslant \mu^*(D_{2n} \cup A_{2n-1})
	= \mu^*(D_{2n}) + \mu^*(A_{2n-1}) \]
	\[ \mu^*(A_{2n}) = \mu^*(D_{2n-1} \cup A_{2n-1}) \geqslant \mu^*(D_{2n-1} \cup A_{2n-2})
	= \mu^*(D_{2n-1}) + \mu^*(A_{2n-2}) \]

	\noindent
	where we can set $A_0 = \emptyset$. It hence follows that (inductively):
	\[ \sum_{k=1}^n \mu^*(D_{2k}) \leqslant \sum_{k=1}^n \mu^*(D_{2k}) + \mu^*(A_1) \leqslant \mu^*(A_{2n+1})
	\leqslant \mu^*(A) < \infty \]
	\[ \sum_{k=1}^n \mu^*(D_{2k-1}) = \sum_{k=1}^n \mu^*(D_{2k-1}) + \mu^*(A_0) \leqslant \mu^*(A_{2n})
	\leqslant \mu^*(A) < \infty \]

	\noindent
	As the above holds for all $n \in \N$, we have in particular that $\sum_{k=1}^\infty \mu^*(D_{2k})$
	and $\sum_{k=1}^\infty \mu^*(D_{2k-1})$ are convergent series. Thus, $\sum_{k=1}^\infty \mu^*(D_k)$ is also a
	convergent series. By subadditivity, we have for all $n \in \N$ that:
	\[ \mu^*(A \setminus F) = \mu^*\left( \bigcup_{k \geqslant 1} A_k \right)
	= \mu^*\left( A_n \cup \bigcup_{k \geqslant n+1} D_k \right)
	\leqslant \mu^*(A_n) +  \sum_{k \geqslant n+1} \mu^*(D_k) \]

	\noindent
	By convergence of the series:
	\[ \mu^*(A \setminus F) \leqslant \liminf_{n \to \infty} \left[ \mu^*(A_n) +  \sum_{k \geqslant n+1} \mu^*(D_k) \right]
	= \liminf_{n \to \infty} \mu^*(A_n) + \liminf_{n \to \infty} \sum_{k \geqslant n+1} \mu^*(D_k)
	= \liminf_{n \to \infty} \mu^*(A_n) \]

	\noindent
	Moreover, $A_n \subseteq A \setminus F$ means $\mu^*(A_n) \leqslant \mu^*(A \setminus F)$ by monotonicity, so in fact:
	\[ \mu^*(A \setminus F) \leqslant \liminf_{n \to \infty} \mu^*(A_n) \leqslant \limsup_{n \to \infty} \mu^*(A_n)
	\leqslant \mu^*(A \setminus F) \]
	\[ \implies \lim_{n \to \infty} \mu^*(A_n) = \mu^*(A \setminus F) \]

	\noindent
	Finally, it follows that:
	\[ \mu^*(A) \geqslant \mu^*(A \cap F) + \lim_{n \to \infty} \mu^*(A_n) = \mu^*(A \cap F) + \mu^*(A \setminus F) \]

	\noindent
	and we hence get that $F$ is $\mu^*$-measurable. As $F$ was an arbitrarily-chosen closed set, it follows that every Borel set is
	$\mu^*$-measurable.
\end{proof}

The following result is therefore immediate (following the fact that $H^d$ is a metric outer measure):

\begin{cor}
	In a metric space $(X, \rho)$, every Borel set is $H^d$-measurable.
\end{cor}

\begin{rem}
	As $\Bor(X)$ is a $\sigma$-algebra and each $A \in \Bor(X)$ is $H^d$-measurable, it follows that
	$(X, \Bor(X), H^d)$ is a measure space (with $H^d$ restricted to $\Bor(X)$).
\end{rem}

The next step is to justify that the Hausdorff measure can indeed be considered a generalization of the Lebesgue
measure (though we will not go into too much detail, as it would take quite some time).
\bigskip

Recall that the Lebesgue outer measure on $\R$ is \textit{translation-invariant}. That is, $m^*(x + A) = m(A)$ for any $A \subseteq \R$
and any $x \in \R$. We also know that $m^*(\lambda A) = |\lambda| m^*(A)$ for $\lambda \in \R$.
The following analogous result holds for the Hausdorff measure:

\begin{lem}\label{isom}
	On a metric space $(X, \rho)$, the Hausdorff outer measure $H^d$ is invariant under isometries of $(X, \rho)$.
	Moreover, for any set $Y$, any $\lambda \in \R$, and functions $f,g : Y \to X$ satisfying $\rho(f(y_1), f(y_2)) \leqslant |\lambda|
	\rho(g(y_1), g(y_2))$ for all $y_1, y_2 \in Y$, we have $H^d(f(B)) = |\lambda|^d H^d(g(B))$ for all $B \subseteq Y$.
\end{lem}
\begin{proof}
	(Adapted from \cite{Fol}).
	\bigskip

	\noindent
	Let $h : X \to X$ be an isometry, so that $\rho(x,y) = \rho(h(x), h(y))$ for all $x,y \in X$.
	In particular, for any $U \subseteq X$, the fact that $h$ is isometric and bijective implies:
	\[ \diam(U) = \sup \{ \rho(x,y) : x,y \in U \} = \sup \{ \rho(h(x), h(y)) : x,y \in U \} \]
	\[ = \sup \{ \rho(h(x), h(y)) : h(x),h(y) \in h(U) \} = \diam(h(U))  \]

	\noindent
	and likewise, $\diam(h^{-1}(U)) = \diam(U)$ (in other words, $\diam$ is invariant under isometries).
	Let $A \subseteq X$. Then for any $\delta > 0$:
	\[ H_\delta^d(h(A)) = \inf \left\{ \sum_{i \geqslant 1} \diam(U_i)^d : h(A) \subseteq \bigcup_{i \geqslant 1} U_i, \; \diam(U_i) < \delta \right\} \]
	\[ = \inf \left\{ \sum_{i \geqslant 1} \diam(h^{-1}(U_i))^d : A \subseteq \bigcup_{i \geqslant 1} h^{-1}(U_i), \; \diam(h^{-1}(U_i)) < \delta \right\}
	= H_\delta^d(A) \]

	\noindent
	and it follows by letting $\delta \to 0$ that $H^d(h(A)) = H^d(A)$, so that $H^d$ is invariant under isometries.
	\bigskip

	\noindent
	For the second statement, we may assume WLOG that $\lambda \neq 0$ (the case where $\lambda = 0$ is clear). Let $B \subseteq Y$
	and let $\delta, \varepsilon > 0$ be arbitrary. Then we can find a cover
	$V_1, V_2, ... \subseteq X$ of $g(B)$ such that $\diam(V_i) < \frac{\delta}{|\lambda|}$ and:
	\[ \sum_{i=1}^\infty \diam(V_i)^d \leqslant H_\delta^d(g(B)) + \frac{\varepsilon}{|\lambda|^d} \]

	\noindent
	Let $U_i = f(g^{-1}(V_i))$. Then we have by construction that $U_1, U_2, ... \subseteq X$ is a cover for $f(A)$,
	and by the assumption we have:
	\[ \diam(U_i) = \sup \{ \rho(x_1, x_2) : x_1, x_2 \in U_i \} = \sup \{ \rho(f(y_1),f(y_2)) : y_1, y_2 \in g^{-1}(V_i) \}\]
	\[ \leqslant |\lambda| \sup \{ \rho(g(y_1),g(y_2)) : y_1, y_2 \in g^{-1}(V_i) \}
	= |\lambda| \sup \{ \rho(x_1,x_2) : x_1, x_2 \in V_i \} = |\lambda| \diam(V_i) \]

	\noindent
	Hence, it follows that:
	\[ H_\delta^d(f(B)) \leqslant \sum_{i=1}^\infty \diam(U_i)^d \leqslant |\lambda|^d \sum_{i=1}^\infty \diam(V_i)^d
	\leqslant |\lambda|^d H_\delta^d(g(B)) + \varepsilon \]

	\noindent
	Now since $\delta$ and $\varepsilon$ were arbitrary, letting $\delta, \varepsilon \to 0$ gives:
	\[ H^d(f(B)) \leqslant |\lambda|^d H^d(g(B)) \]

	\noindent
	which completes the proof.
\end{proof}

The following result justifies that the Hausdorff measure can be considered a generalization of the Lebesgue measure:

\begin{prop}\label{leb}
	For $\R^n$ equipped with the usual Euclidean metric and $m_n$ the Lebesgue measure on $\R^n$, there exists
	a constant $\gamma_n > 0$ for which $H^n = \gamma_n m_n$.
\end{prop}
\begin{proof}
	We will \textit{sketch} the general idea of the proof. Consider the unit cube $Q = [0,1]^n$ and define $\gamma_n = H^n(Q)$.
	We can show that $\gamma_n < \infty$ by choosing $\delta > 0$, writing $Q$ as a union of cubes of length
	$< \frac{\delta}{\sqrt{n}}$, and verifying that the sum of diameters of the cubes is $n^{n/2}$. It will follow that
	$\gamma_n = H^n(Q) < \infty$ (as $\delta$ doesn't depend on $n$).
	\bigskip

	\noindent
	We can similarly show that $\gamma_n > 0$ by choosing an arbitrary countable cover $U_1, U_2, ...$ of $Q$ and covering
	each $U_i$ by a cube $Q_i$ of side length $\diam(U_i)$. Monotonicity will then give
	that $m_n(Q) = 1 \leqslant \sum_{i=1}^\infty m_n(Q_i) = \sum_{i=1}^\infty (\diam(U_i))^n$, and taking the infimum over all such
	covers leads to $1 \leqslant H^n(Q)$ (so that in particular, $\gamma_n = H^n(Q) > 0$).
	\bigskip

	\noindent
	Finally, we can make use of the uniqueness of the Carath\'{e}odory theorem
	in the construction of $m_n$ to show that $H^n(B) = \gamma_n m_n(B)$ for all open \textit{boxes} $B \subseteq \R^n$.
	This is done by approximating the boxes by ``almost disjoint'' cubes (as done in \textbf{A1 Q5 (2)}) and using \textbf{Lemma \ref{isom}}
	to deal with the scaling and translating of each cube to $[0,1]^n$. It will then follow that in fact $H^n(A) = \gamma_n m_n(A)$
	for all Lebesgue-measurable $A \subseteq \R^n$.
\end{proof}

\begin{rem}\label{vol_rem}
	The proof doesn't require the computation of the value of $\gamma_n$, but only that it's a (strictly) positive finite value. It turns out that
	$\gamma_n = \frac{\vol(B_1^n)}{2^n}$ where $\vol(B_1^n)$ is the volume of the $n$-ball $B_1^n \subseteq \R^n$ of radius 1
	(diameter 2). In particular, $\gamma_1 = 1$ so that $H^1$ is exactly the 1-dimensional Lebesgue measure $m_1$. We will prove this
	latter fact formally (as it will help later).
\end{rem}

\begin{lem}\label{gamma_1}
	$\gamma_1 = 1$. Hence, $H^1 = m_1$ (that is, the 1-dimensional Hausdorff measure on $\R$ is exactly the Lebesgue measure on $\R$).
\end{lem}
\begin{proof}
	Since $H^1([0,1]) = \gamma_1 m_1([0,1]) = \gamma_1$, we must show that $H^1([0,1]) = 1$.
	Indeed, let $\delta > 0$ and choose $N \in \N$ for which $\frac{1}{N} < \delta$. Then for each $1 \leqslant n \leqslant N$, set
	$U_n = [\frac{n-1}{N}, \frac{n}{N}]$ so that $U_1, ..., U_N$ is a cover for $[0,1]$ with $\diam(U_n) = \frac{1}{N} < \delta$.
	Then:
	\[ H_\delta^1([0,1]) \leqslant \sum_{n=1}^N \diam(U_n) \leqslant \sum_{n=1}^N \frac{1}{N} = 1 \]

	\noindent
	and so letting $\delta \to 0$, we have $H^1([0,1]) \leqslant 1$. On the other hand, let $\delta > 0$ and
	choose any cover $U_1, U_2, ... \subseteq \R$ for $[0,1]$ satisfying $\diam(U_i) < \delta$. We can see that $\diam(U_i) \geqslant m_1^*(U_i)$
	(where $m_1^*$ is the Lebesgue \textit{outer} measure)
	because $\diam(U_i) < \delta$ implies we can find $\alpha_i = \inf(U_i)$ and $\beta_i = \sup(U_i)$ so that
	$\diam(U_i) = \beta_i - \alpha_i$ and $J_i = [\alpha, \beta] \supseteq U_i$ so that $\diam(U_i) = \beta_i - \alpha_i
	= m_1(J_i) \geqslant m_1^*(U_i)$ by monotonicity. It then follows (by subadditivity and monotonicity) that:
	\[ \sum_{i=1}^\infty \diam(U_i) \geqslant \sum_{i=1}^\infty m_1^*(U_i) \geqslant m_1^*\left( \bigcup_{i=1}^\infty U_i \right)
	\geqslant m_1([0,1]) = 1 \]

	\noindent
	so that (since our choice of cover was arbitrary) $H_\delta^1([0,1]) \geqslant 1$. Letting $\delta \to 0$ then gives
	$H^1([0,1]) \geqslant 1$, and thus $H^1([0,1]) = \gamma_1 = 1$.
\end{proof}

\section{Hausdorff dimension}

Before defining the Hausdorff dimension of a set, we first need the following lemma:

\begin{lem}\label{dim}
	Let $(X, \rho)$ be a metric space, let $E \subseteq X$, and let $\alpha, \beta \in \R$ with $0 < \alpha < \beta$.
	If $H^\alpha(E) < \infty$, then $H^\beta(E) = 0$.
\end{lem}
\begin{proof}
	(Adapted from \cite{Roy}).
	\bigskip

	\noindent
	Let $\delta > 0$. As $H^\alpha(E) < \infty$, we can choose $U_1, U_2, ... \subseteq X$ for which $E \subseteq \bigcup_{i \geqslant 1} U_i$,
	each $\diam(U_i) < \delta$, and:
	\[ \sum_{i \geqslant 1} \diam(U_i)^\alpha < H_\delta^\alpha(E) + 1  \]

	\noindent
	We then see that as $\alpha < \beta$ and $\diam(U_i) < \delta$:
	\[ H_\delta^\beta(E) \leqslant \sum_{i \geqslant 1} \diam(U_i)^\beta = \sum_{i \geqslant 1} \diam(U_i)^{\beta - \alpha + \alpha}
	\leqslant \delta^{\beta - \alpha} \sum_{i \geqslant 1} \diam(U_i)^\alpha
	\leqslant \delta^{\beta - \alpha} \left( H_\delta^\alpha(E) + 1 \right) \]

	\noindent
	Moreover, since $H_\delta^\alpha(E)$ is increasing as $\delta \to 0$ (as proved in \textbf{A3}):
	\[ H_\delta^\beta(E) \leqslant \delta^{\beta - \alpha} \left( H_\delta^\alpha(E) + 1 \right) \leqslant \delta^{\beta - \alpha}
	\left( H^\alpha(E) + 1 \right) \]

	\noindent
	Finally, we have:
	\[ H^\beta(E) = \lim_{\delta \to 0} H_\delta^\beta(E) \leqslant \lim_{\delta \to 0} \delta^{\beta - \alpha} \left( H^\alpha(E) + 1 \right)
	= 0 \]

	\noindent
	which forces $H^\beta(E) = 0$.
\end{proof}

The lemma makes the following definition well-defined:

\begin{defn}
	Let $(X, \rho)$ be a metric space and let $E \subseteq X$. The \textbf{Hausdorff dimension} $\dim_H(E)$ of $E$ is defined to be:
	\[ \dim_H(E) = \inf \{ \alpha \geqslant 0 : H^\alpha(E) = 0 \} \]
\end{defn}

\paragraph{}
Using \textbf{Lemma \ref{dim}} and \textbf{Proposition \ref{leb}}, we can compute the Hausdorff dimensions of certain shapes
(subsets of $\R^n$). The following are examples in $\R^2$.

\begin{exa}
	Consider the closed unit disk $\D = \{ x \in \R^2 : \norm{x}_2 \leqslant 1 \}$.
	We claim that its Hausdorff dimension $\dim_H(\D) = 2$. By using \textbf{Proposition \ref{leb}}, we have that:
	\[ H^2(\D) = \gamma_2 m_2(\D) = \gamma_2 \cdot \pi > 0 \]

	\noindent
	where $m_2(\D) = \pi$ is just the $\R^2$-Lebesgue measure (area) of $\D$. Then \textbf{Lemma \ref{dim}} tells us that
	$H^\alpha(\D) = 0$ for all $\alpha > 2$, and so:
	\[ \dim_H(\D) = \inf\{ \alpha \geqslant 0 : H^\alpha(\D) = 0 \} = 2 \]

	\noindent
	We used \textbf{Proposition \ref{leb}} to prove that $H^2(\D) > 0$, but we could also have done an infimum argument
	to do it. The former method saves us time, however.
\end{exa}

It can be more difficult to compute Hausdorff dimensions of more general submanifolds of $\R^n$, as
\textbf{Proposition \ref{leb}} would not guarantee that $H^n$ gives a non-zero value. In the next example, we compute
the Hausdorff dimension of the submanifold $S^1$ (the unit circle in $\R^2$), and the following lemma proves useful to do so:

\begin{lem}\label{conn}
	Let $(X, \rho)$ be a metric space and let $A \subseteq X$ be a \textit{connected} set. Then $H^1(A) \geqslant \diam(A)$.
\end{lem}
\begin{proof}
	(Adapted from \cite{Sem})
	\bigskip

	\noindent
	For each $a \in X$, consider the map $d_a : X \to \R$ defined by $d_a(x) = \rho(a,x)$. We can see that for $x,y \in X$
	(assuming WLOG $\rho(a,x) \geqslant \rho(a,y)$):
	\[ |d_a(x) - d_a(y)| = |\rho(a,x) - \rho(a,y)| = \rho(a,x) - \rho(a,y) \leqslant \rho(a,y) + \rho(x,y) - \rho(a,y) = \rho(x,y) \]

	\noindent
	where the inequality comes from the triangle inequality on $\rho$. This shows that $d_a$ is 1-Lipschitz, thus (in particular) continuous.
	We hence get that:
	\[ \diam(d_a(U)) = \sup\{ |d_a(x)-d_a(y)| : x,y \in U \} \leqslant \sup\{ \rho(x,y) : x,y \in d_a(U) \} = \diam(U) \]

	\noindent
	for any $U \subseteq X$, which implies for any $\delta > 0$:
	\[ H_\delta^1(d_a(A)) = \inf\left\{ \sum_{i \geqslant 1} \diam(J_i) : d_a(A) \subseteq \bigcup_{i \geqslant 1} J_i, \; \diam(J_i) < \delta \right\} \]
	\[ \leqslant \inf\left\{ \sum_{i \geqslant 1} \diam(d_a(U_i)) : A \subseteq \bigcup_{i \geqslant 1} U_i, \; \diam(d_a(U_i)) \leqslant
	\diam(U_i) < \delta \right\} \]
	\[ \leqslant \inf\left\{ \sum_{i \geqslant 1} \diam(U_i) : A \subseteq \bigcup_{i \geqslant 1} U_i, \; \diam(U_i) < \delta \right\}
	= H_\delta^1(A) \]

	\noindent
	so that $H^1(d_a(A)) \leqslant H^1(A)$ for all $a \in X$. Finally since $d_a$ is continuous and $A$ is connected, it follows that
	$d_a(A) \subseteq \R$ is connected (i.e. an interval), which means $H^1(d_a(A)) = \gamma_1 m_1(d_a(A)) = \gamma_1 \diam(d_a(A))$.
	By \textbf{Lemma \ref{gamma_1}}, $\gamma_1 = 1$ and so $H^1(d_a(A)) = \diam(d_a(A)) \leqslant H^1(A)$ for all $a \in X$.
	Finally:
	\[ \sup_{a \in A} \diam(d_a(A)) = \sup_{a \in X} \sup \{ |d_a(x)-d_a(y)| : x,y \in A \}
	= \sup_{a \in X} \sup \{ |\rho(x,a)-\rho(y,a)| : x,y \in A \} \]
	\[ = \sup_{a \in X} \sup \{ \rho(x,a)-\rho(y,a) : x,y \in A \}
	\geqslant \sup \{ \rho(x,y)-\rho(y,y) : x,y \in A \}\]
	\[ = \sup \{ \rho(x,y) : x,y \in A \} = \diam(A)   \]

	\noindent
	so that $\diam(A) \leqslant \sup_{a \in A} \diam(d_a(A)) \leqslant H^1(A)$, which completes the proof.
\end{proof}

\begin{exa}
	Consider the unit circle $S^1 = \{ x \in \R^2 : \norm{x}_2 = 1 \}$.
	We claim that its Hausdorff dimension $\dim_H(S^1) = 1$. We know that $S^1$ has $\R^2$-Lebesgue measure zero, so that:
	\[ H^2(S^1) = \gamma_2 m_2(S^2) = \gamma_2 \cdot 0 = 0 \]

	\noindent
	But this only tells us that $\dim_H(S^1) \leqslant 2$. To show that $\dim_H(S^1) = 1$, we observe that
	$S^1 \subseteq \R^2$ is \textit{connected}, and so by \textbf{Lemma \ref{conn}}:
	\[ H^1(S^1) \geqslant \diam(S^1) = 2 > 0 \]

	\noindent
	so that $\dim_H(S^1) = 1$ (by \textbf{Lemma \ref{dim}}).
\end{exa}

\newpage

\begin{thebibliography}{MMMMM}
\bibitem[Folland]{Fol} Folland, G. B. (1999). \textit{Real analysis: Modern techniques and their applications}. New York: Wiley.
\bibitem[Royden]{Roy} Royden, H. L. \& Fitzpatrick, P. (2010). \textit{Real analysis}. Boston: Prentice Hall.
\bibitem[Semmes]{Sem} Semmes, S. (2010). \textit{Some elementary aspects of Hausdorff measure and dimension}. arXiv:1008.2637v1.
\end{thebibliography}

\bibliographystyle{plain}
\bibliography{template}

\end{document}


\lab{Policy Function Iteration}{Policy Function Iteration}
\objective{Iterative methods can be powerful ways to solve dynamic optimization problems without computing the exact solution.
Often we can iterate very quickly to the true solution, or at least within some $\epsilon$ error of the solution.
These methods are significantly faster than computing the exact solution using dynamic programming.
We demonstrate two iterative methods, value iteration (VI) and policy iteration (PI), and use them to solve a deterministic Markov decision process.
}

\section*{Dynamic Optimization}

Many dynamic optimization problem take the form of a \emph{Markov decision process}.
They are formulated as follows.

$\mathbb{T}$ is a set of discrete time periods.
In this lab, $\mathbb{T} = {0,1,\ldots, T}$.
$S$ is the set of possible states.
The set of allowable actions for each state $s$ is $A_s$.
$s_{t+1}=g(s_t,a_t)$ is a transition function that describes how the state changes, based on the previous state and action.
The reward $u_t(s,a)$ is the reward for taking action $a$ while in state $s$ at time $t$.
If the Markov process is stochastic, $p_t(s,a)$ is probability of taking action $a$ at time $t$ while in state $s$.
A deterministic Markov process has $p_t(s,a) = 1$ $\forall s,a$.
The time discount factor $\beta \in [0,1]$ determines how much less a reward is worth in the future.
Then the dynamic optimization problem is

\begin{align}
\label{eq:policyiter-dynopt1}
\max_a  & \sum_{t=0}^T \beta^t u(s_t,a_t) \\
\mbox{subject to } & s_{t+1}= g(s_t,a_t)\ \forall t.
\end{align}

The cake eating problem described in the previous lab follows this format where $S$ consists of the possible amounts of remaining cake ($\frac{i}{W}$), $c_t$ is the amount of cake we can eat, and the amount of cake remaining $s_{t+1}=g(s_t,a_t)$ is $w_t-c_t$, the amount of cake we have, $w_i$, minus the cake we eat, $c_t$.

\subsection*{Moving on a Grid}

Now consider an $N \times N$ grid.
Assume that a robot moves around the grid, one space at a time, until it reaches the lower right hand corner and stops.
Each square is a state, $S = \{0, 1, \ldots, N^2-1\}$, and the set of actions is $\{Left, Down, Right, Up\}$.
For this lab, $Left = 0$, $Down = 1$, $Right = 2$, and $Up = 3$.

\begin{warn}
It is important to remember that the actions do not correspond to the states the robot is in after the action.
When the robot is in state $0$ and takes action $1$, he is then in state $2$.
\end{warn}

$A_s$ is the set of actions that keep the robot on the grid.
If the robot is in the top left hand corner, the only allowed actions are $Down$ and $Right$ so $A_0 = \{1,2\}$.
The transition function $g(s_t,a_t) = s_{t+1}$ can be explicitly defined for each $s, a$ where $s_{t_1}$ is the new state after moving.
For this lab, we define a dictionary $P$ to represent the decision process.
P[state][action]=[(prob, nextstate, reward, is\_terminal),...]
This dictionary contains all of the information about the states, actions, probabilities, and rewards.
The final entry in the tuple, is\_terminal, indicates if the new state is a stopping point.

Let $N=2$ and label the squares as displayed below.
In this example, we define the reward to be $-1$ if the robot moves into 2, $-1$ if the robot moves into $0$ from $1$, and $1$ when it reaches the end, $3$.
This simplifies to $u_t(s,a) = u(a)$ except when moving into state $0$.
Similarly, $p_t(s,a) = p(a) = 1$.

\begin{center}
\begin{tabular}{|c|c|}
\hline
0 & 1 \\ \hline
\cellcolor{red!20}2 & \cellcolor{green!20}3 \\ \hline
\end{tabular}
\end{center}

%\begin{align*}
%S &= \{0, 1, 2, 3\}\\
%A_0 &= \{Right, Down\} &A_1 &= \{Left, Up\} &A_2 &= \{Up, Right\} & A_3 &= \{\} \\
%g(0,Right) &= 1 &g(1,Left) &= 0 &g(2,Up) &= 0\\
%g(0,Down) &= 2 &g(1,Down) &= 3 &g(2,Right) &= 3 \\
%u(0) &= 0 &u(1) &= 0 &u(2) &= -1 &u(3) &= 1 \\
%p(0) &= .75 &p(1) &= .25 &p(2) &= .75 &p(3) &= .25
%\end{align*}

All of this information is encapsulated in $P$.
We define $P[state][action]$ for all states and actions, even if they are not possible.
This simplifies coding the algorithm but is not necessary.

\begin{center}
\begin{tabular}{llll}
\li{P[0][0] = [(0, 0, 0, False)]}
    & \li{P[2][0] = [(0, 2, -1, False)]}\\
\li{P[0][1] = [(1, 2, -1, False)]}
    & \li{P[2][1] = [(0, 2, -1, False)]}\\
\li{P[0][2] = [(1, 1, 0, False)]}
    & \li{P[2][2] = [(1, 3, 1, True)]}\\
\li{P[0][3] = [(0, 0, 0, False)]}
    & \li{P[2][3] = [(1, 0, 0, False)]}\\
\li{P[1][0] = [(1, 0, -1, False)]}
    &\li{P[3][0] = [(0, 0, 0, True)]} \\
\li{P[1][1] = [(1, 3, 1, True)]}
    &\li{P[3][1] = [(0, 0, 0, True)]}\\
\li{P[1][2] = [(0, 0, 0, False)]}
    &\li{P[3][2] = [(0, 0, 0, True)]}\\
\li{P[1][3] = [(0, 0, 0, False)]}
    &\li{P[3][3] = [(0, 0, 1, True)]}
\end{tabular}
\end{center}

We define the \emph{value function} $V(s)$ to be the maximum possible reward of starting in state $s$.
Then using Bellman's optimality equation,
\begin{equation}
\label{eq:policyiter-val-func}
V(s) = \max_{a \in A_s} \{\Sigma[p(a) * \left( u(a) + \beta V(a)\right)]\}.
\end{equation}

The summation occurs when $p(s,a)<1$ so $P[s][a]$ consists of more than one tuple.
For example, if the robot is in the top left corner and moves right, we could have that the probability the robot actually moves right is $.5$.
In this case, $P[0][2] = [(.5, 1, 0, False), (.5, 2, -1, False)]$.
This will occur later in the lab.
We also observe that in this scenario, $V(a)$ is $V(s')$ where $s'$ is the state after taking action $a$ in state $s$.

\section*{Value Iteration}

In the previous lab, we used dynamic programming to solve for the value function.
This was a recursive method where we calculated all possible values for each state and time period.
\emph{Value iteration} is another algorithm that solves the value function by taking an initial value function and calculating a new value function iteratively.
Since we are not calculating all possible values, it is typically faster than dynamic programming.

\subsection*{Convergence of Value Iteration}

A function $f$ that is a contraction mapping has a \emph{fixed point} $p$ such that $f(p) = p$.
Blackwell's contraction theorem can be used to show that Bellman's equation is a ``fixed point'' (it actually acts more like a fixed function in this case)
for an operator $T: L^{\infty}(X;\mathbb{R}) \to L^{\infty}(X;\mathbb{R})$ where $L^{\infty}(X;\mathbb{R})$ is the set of all bounded functions:
\begin{equation}
\label{eq:policyiter-blackwell}
[T(f)](s) = \max_{a \in A_s} \{\Sigma [p(a) * \left( u(a) + \beta f(a)\right)]\}
\end{equation}
It can be shown that \ref{eq:policyiter-dynopt1} is the fixed ``point'' of our operator $T$.
A result of contraction mappings is that there exists a unique solution to \ref{eq:policyiter-blackwell}.

\begin{equation}
\label{eq:policyiter-val-iteration}
V_{k+1}(s_i) = [T(V_k)](s_i) = \max_{a \in A_s} \{\Sigma[p(a) *\left( u(a) + \beta V_k(a)\right)]\}
\end{equation}

where an initial guess for $V_0(s)$ is used.
As $k \to \infty$, it is guaranteed that $(V_k(s)) \to V^*(s)$.
Because of the contraction mapping, if $V_{k+1}(s) = V_k(s) \, \, \forall \, s$, we have found the true value function, $V^*(s)$.
Using this information, we define the value iteration algorithm to find $V^*$:

\begin{algorithm}[H]
\begin{algorithmic}[1]
\Procedure{Value Iteration Function}{$P, S, A, \beta, \varepsilon$, maxiter}
    \State $V_0 \gets [V_0(s_0),V_0(s_1),\ldots,V_0(s_N)] $
     \Comment{Common choice is $V_0(s_i)=u(s_i)$}
    \For{$i=1,2,\dots,$\ \li{maxiter}}
        \Comment Iterate only \li{maxiter} times at most.
        \For{$s \in S$}
             \State $V_{k+1}(s) = \max_{a \in A_s}\{\Sigma[p(a)*(u(a) + \beta*V_k(a))]\}$
        \EndFor
        \If{ $||V_{k+1} - V_k|| < \varepsilon$}
            \State \texttt{break}
             \Comment{Stop iterating if the approximation stops changing enough.}
        \EndIf
    \EndFor
    \State \pseudoli{return} $V_k$
\EndProcedure
\end{algorithmic}
\caption{Value Function Iteration}
\label{alg:ValueIteration}
\end{algorithm}

Let $V_0 = [0,0,0,0]$ and $\beta = 1$.
We calculate $V_1(s)$ from the example above.
\begin{align*}
V_1(0) &= \max_{a \in A_0} \{\Sigma[p(a)*(u(a)+V_0(a))]\} \\
&= \max \{p(1)*(u(1)+V_0(1)), p(2)*(u(2)+V_0(2)))\} \\
&= \max \{1(-1+0), 1(0+0)\} \\
&= \max \{-1,0\} \\
&= 0\\
V_1(1) &= \max \{p(0)*(u(0)+V_0(0)), p(2)*(u(2)+V_0(2))\}\\
&=\max \{1(0+0), 1(1+0)\} \\
&= 1
\end{align*}

Calculating $V_1(2)$ and $V_1(3)$ gives $V_1 = [0, 1, 1, 0]$.
Repeating the process, $V_2 = [1, 1, 1, 0]$, which is the solution.
It means that maximum reward the robot can achieve by starting on square $i$ is $V_2[i]$.

Most iterative algorithms have a \li{maxiter} parameter that will terminate the algorithm after \li{maxiter} iterations regardless of whether or not it has converged.
This is because even though we have guaranteed convergence, we might have a convergence rate that is too slow to be useful.
However, generally this algorithm will converge much faster than computing the true value function using dynamic programming.

\begin{problem}
\label{prob:policyiter-value1}
Write a function called \li{value_iteration()} that will accept a dictionary $P$ representing the decision process, the number of states, the number of actions, a discount factor $\beta \in (0,1)$, defaulting to $1$,
the tolerance amount $\epsilon$ defaulting to $1e^-8$, and the maximum number of iterations \li{maxiter} defaulting to $3000$.
Perform value iteration until $\|V_{k+1} - V_{k}\| < \epsilon$ or $k > $ \li{max\_iter}.
Return the final vector representing $V^*$ and the number of iterations.
Test your code on the example given above.
\end{problem}

\subsection*{Calculating the Policy}

While knowing the maximum expected value is helpful, it is usually more important to know the policy that generates the most value.
Value Iteration tells the robot what reward he can expect, but not how to get it.
The policy vector, $\mathbf{c}$, is found by using the policy function: $\pi : \mathbb{R} \to \mathbb{R}$.
$\pi(s)$ is the action we should take while in state $s$ to maximize reward.
We can modify the Bellman equation using $V^*(s)$ to find $\pi$:
\begin{equation}
\label{eq:policyiter-pol_func}
\pi(s) = \argmax_{a \in A_s} \{\Sigma [p(a)*(u(a) + \beta*V^*(s))]\}
\end{equation}

Using value iteration, we found $V^*  = [1, 1, 1, 0]$ in the example above.
We find $\pi(0)$ by
\begin{align*}
\pi(0) &=\argmax_{1,2} \{p(1)*(u(1) + V^*(1)), p(2)*(u(2)+V^*(2))\} \\
&= \argmax \{1*(-1+1), 1*(0+1)\} \\
&= \argmax\{0, 1\}\\
&= 2.
\end{align*}

So when the robot is in state $0$, he should take action $2$, moving $Right$.
This avoids the $-1$ penalty for moving $Down$ into square $2$.
Similarly,
\begin{align*}
\pi(1) &= \argmax_{0,1} \{1*(0+1), 1*(1+1)\} = \argmax \{1, 2\} = 1.
\end{align*}
The policy corresponding to the optimal reward is $[2,1,2,0]$.
The robot should move to square $3$ if possible, avoiding $2$ because it has a negative reward.
Since $3$ is terminal, it does not matter what $\pi(3)$ is.
We set it to $0$ for convenience.

\begin{problem}
\label{prob:policyiter-value2}
Write a function called \li{extract_policy()} that will accept a dictionary $P$ representing the decision process, the number of states, the number of actions, an array representing the value function, and a discount factor $\beta \in (0,1)$, defaulting to $1$.
Return the policy vector corresponding to $V^*$.
Test your code on the example with $\beta = 1$.
\end{problem}

\section*{Policy Iteration}
For dynamic programming problems, it can be shown that value function iteration converges relative to the discount factor $\beta$.
As $\beta$ approaches $1$, the number of iterations increases dramatically.
As mentioned earlier $\beta$ is usually close to $1$, which means this algorithm can converge slowly.
In value iteration, we used an initial guess for the value function, $V_0$ and used (\ref{eq:policyiter-dynopt1}) to iterate towards the true value function.
Once we achieved a good enough approximation for $V^*$, we recovered the true policy function $\pi$.
Instead of iterating on our value function, we can instead make an initial guess for the policy function, $\pi_0$, and use this to iterate toward the true policy function, called policy iteration.
We do so by taking advantage of the definition of the value function, where we assume that our policy function yields the most optimal result.

That is, given a specific policy function $\pi_k$, we can modify \eqref{eq:policyiter-dynopt1} by assuming that the policy function is the optimal choice.
This process, called \emph{policy evaluation}, evaluates the value function for a given policy.

\begin{equation}
\label{eq:policyiter-val_from_policy}
V_k(s) = \max_{a \in [A_s]} \{\Sigma([p(a)*(u(a) + \beta V^*(s))]\} = \max_{a \in [A_s]} \Sigma([p(\pi(s))*(u(\pi(s)) + \beta *V^*(\pi(s)))]\}
\end{equation}
The last equality occurs because in state $s$, the robot should choose the action that maximizes reward, which is $\pi(s)$ by definition.

\begin{problem}
\label{prob:policyiter-value3}
Write a function called \li{compute_policy_v()} that accepts a dictionary $P$ representing the decision process, the number of states, the number of actions, an array representing a policy, a discount factor $\beta \in (0,1)$ defaulting to $1$, and a tolerance amount $\epsilon$ defaulting to $1e^-8$,.
Return the value function corresponding to the policy.
Test your code on the policy vector generated from \li{extract_policy} for the example.
The result should be the same value function array from \li{value_iteration}.
\end{problem}

Now that we have the value function for our policy, we can take the value function and find a better policy.
Called policy improvement, this step is the same method used in value iteration to find the policy.

Thus, given an initial guess for our policy function, $\pi_0$, we calculate the corresponding value function using (\ref{eq:policyiter-val_from_policy}), and then use \eqref{eq:policyiter-pol_func} to improve our policy function.
The algorithm for policy function iteration can be summarized as follows:

\begin{algorithm}[H]
\begin{algorithmic}[1]
\Procedure{Policy Iteration Function}{$P, nS, nA, \beta, tol,$ maxiter}
    \State $\pi_0 \gets [\pi_0(w_0),\pi_0(w_1),\ldots,\pi_0(w_N)] $
     \Comment{Common choice is $\pi_0(w_i)=w_{i-1}$ with $\pi_0(0)=0$}
    \For{$i=1,2,\dots,$\ \li{maxiter}}
        \Comment Iterate only \li{maxiter} times at most.
        \For{$s \in S$}
        \Comment{Policy evaluation}
            \State $V_{k}(s) = \max_{a \in A_s} \{\Sigma[p(\pi(s))*(u(\pi(s)) + \beta*V^*(\pi(s)))]\}$
            \Comment{compute\_policy\_v.}
        \EndFor
        \For{$s \in S$}
        \Comment{Policy improvement.}
            \State $\pi_{k+1}(s) =\argmax_{a \in A_s} \{\Sigma[p(a)*(u(a) + \beta*V_k*(s))]\}$
            \Comment{extract\_policy.}
        \EndFor
        \If{ $||\pi_{k+1} - \pi_k|| < \varepsilon$}
            \State \texttt{break}
             \Comment{Stop iterating if the policy doesn't change enough.}
        \EndIf
    \EndFor
    \State \pseudoli{return} $V_k, \pi_{k+1}$
\EndProcedure
\end{algorithmic}
\caption{Policy Iteration}
\label{alg:PolicyIteration}
\end{algorithm}

\begin{problem}
\label{prob:policyiter-value4}
Write a function called \li{policy_iteration()} that will accept a dictionary $P$ representing the decision process, the number of states, the number of actions, a discount factor $\beta \in (0,1)$, defaulting to $1$, the tolerance amount $\epsilon$ defaulting to $1e^-8$, and the maximum number of iterations \li{maxiter} defaulting to $3000$.
Perform policy iteration until $\|\pi_{k+1} - \pi_{k}\| < \epsilon$ or $k > $ \li{maxiter}.
Return the final vector representing $V_k$, the optimal policy $\pi_k$, and the number of iterations.
Test your code on the example given above and compare your answers to the results from \ref{prob:policyiter-value1} and \ref{prob:policyiter-value2}.
\end{problem}

\section*{The Frozen Lake Problem}
For the rest of the lab, we will be using the OpenAi Gym environments.
It can be installed using conda or pip.
\begin{lstlisting}[language=Bash]
$ pip install gym[all]
\end{lstlisting}

% TODO: INSERT Figure

In the Frozen Lake problem, you and your friends tossed a frisbee onto a mostly frozen lake.
The lake is divided into an $N \times N$ grid where the top left hand corner is the start, the bottom right hand corner is the end, and the other squares are either frozen or holes.
To retrieve the frisbee, you must successfully navigate around the melted ice without falling.
The possible actions are left, right, up, and down.
Since ice is slippery, you won't always move in the intended direction.
If you fall, your reward is $0$.
If you succeed, your reward is $1$.
There are two scenarios: $N=4$ and $N=8$.
To run the $4\times 4$ scenario, use \li{env_name='FrozenLake-v0'}.
For the $8\times 8$ scenario, use \li{env_name='FrozenLake8x8-v0'}.

\subsection*{Using Gym}
To use gym, we import it and create an environment based on the built-in gym environment.
The FrozenLake environment has $3$ important attributes, $P$, $nS,$ and $nA$.
$P$ is a dictionary where $P[s][a]=[(probability, nextstate, reward, is\_terminal)...]$.
Notice that this is the same $P$ we used in the previous problems.
$nS$ and $nA$ are the number of states and actions respectively.

We can calculate the optimal policy using value iteration or policy iteration.
Since the ice is slippery, this policy will not always result in a reward of $1$.
The gym environments have built-in functions that allow us to simulate each step of the environment.
Before running a scenario in gym, always put it in the starting state by calling \li{env.reset()}.
To simulate moving, call \li{env.step}.

\begin{lstlisting}
$ import gym
$ from gym import wrappers
$ env_name  = 'FrozenLake-v0'
$ env = gym.make(env_name).env
$ number_of_states = env.nS
$ obs = env.reset()
$ obs, reward, done, _ = env.step(int(policy[obs]))
\end{lstlisting}

The step function returns four values: observation, reward, done, info.
The observation is an environment-specific object representing the observation of the environment.
For FrozenLake, this is the current state.
When we step, or take an action, we get a new observation, or state, as well as the reward for taking that action.
If we fall into a hole or reach the frisbee, the simulation is over so we are done.
The info value is a dictionary of diagnostic information.
It will not be used in this lab.

\begin{problem}
\label{prob:policyiter-value5}
Write a function that runs \li{value_iteration} and \li{policy_iteration} on FrozenLake.
It should accept a boolean \li{basic_case} defaulting to \li{True} that indicates whether to run the $4\times 4$ or $8\times 8$ scenario and an integer $n$ defaulting to $1000$ that indicates how many times to run the simulation.
Calculate the value function and policy for the environment using both value iteration and policy iteration.
Return the policies generated by value iteration and the policy and value function generated by policy iteration.
Set the mean total rewards to $0$ and return them as well.
\end{problem}

\begin{problem}
\label{prob:policyiter-value6}
Write a function \li{run_simulation()} that takes in the environment \li{env}, a policy \li{policy}, a discount factor $\beta$ defaulting to $1$, and a boolean \li{render} defaulting to False.
Calculate the total reward for the policy for one simulation using \li{env.reset} and \li{env.step()}.
Stop the simulation when $done$ is $True$.
(Hint: When calculating reward, use $\beta^k$.)

Modify \li{frozen_lake()} to run \li{run_simulation} for both the value iteration policy and the policy iteration policy $M$ times.
Return the policy generated by value iteration, the mean total reward for the policy generated by value iteration, the policy generated by policy iteration, and the mean total reward for the policy generated by policy iteration.
\end{problem}

\begin{comment}
\newpage

%
%
%%\begin{equation}
%%\label{eq:policyiter-bellman1}
%%V(w_i) = \max_{c \in [0,w_i]} u(c) + \beta V(w_i-c)
%%\end{equation}
%%It is common to make the substitution $w_i' = w_i - c$ in \eqref{eq:policyiter-bellman1} to obtain the %following.
%\begin{equation}
%\label{eq:policyiter-bellman2}
%V(w_i) = \max_{w_j \in [0,w_i]} u(w_i-w_j) + \beta V(w_j).
%\end{equation}
%
%Notice that \eqref{eq:policyiter-bellman2} is similar to the value function given in the dynamic programming lab.
%%Note that \eqref{eq:policyiter-bellman1} requires finding $c$, how much cake to consume in each period.
%%The variable $c$ is commonly described as the \emph{action} to take when we have $w_i$ cake.
%
%
%\subsection*{Convergence of Value Iteration}
%
%A function $f$ that is a contraction mapping has a fixed point $p$ such that $f(p) = p$.
%Blackwell's contraction theorem can be used to show that Bellman's equation is a ``fixed point'' (it actually acts more like a fixed function in this case)
%for an operator $T: L^{\infty}(X;\mathbb{R}) \to L^{\infty}(X;\mathbb{R})$ where $L^{\infty}(X;\mathbb{R})$ is the set of all bounded functions:
%\begin{equation}
%\label{eq:policyiter-blackwell}
%[T(f)](w) = \max_{w_j \in [0,w]} u(w-w_j) + \beta f(w_j).
%\end{equation}
%It can be shown that (\ref{eq:policyiter-bellman2}) is the "fixed point" of our operator $T$.
%A result of contraction mappings is that there exists a unique solution to (\ref{eq:policyiter-blackwell}).
%
%A powerful property of contraction mappings is that the fixed point can be found by applying the function repeatedly to some initial point in our domain.
%That is, if $f:X \to X$ is our contraction mapping, with fixed point $p \in X$, we can find $p$ by randomly choosing $x \in X$ and repeatedly applying the function $f$ to our point.
%This can be expressed as
%\begin{align*}
%f^{M}(x) = f \circ f \cdots \circ f(x) = p
%\end{align*}
%for some M.  Here M is dependent on both the initial guess $x$ and the discount factor of the contraction mapping, which correspond to $V_0$ and $\beta$ in dynamic optimization.
%
%In the case of dynamic optimization, this implies that we can converge to the true value function $V^*(w)$ by using the following equation:
%\begin{equation}
%\label{eq:policyiter-val_iteration}
%V_{k+1}(w_i) = [T(V_k)](w_i) = \max_{w_j \in [0,w_i]} u(w_i-w_j) + \beta V_k(w_j) \, \, \forall \, w_i,
%\end{equation}
%
%where an initial guess for $V_0(w)$ is used.
%As $k \to \infty$, it is guaranteed that $(V_k(w)) \to V^*(w)$.
%Because of the contraction mapping, if $V_{k+1}(w) = V_k(w) \, \, \forall \, w$, we have found the true value function, $V^*(w)$.
%Using this information, we define the value iteration algorithm to find $V^*$:
%
%
%
%
%\begin{algorithm}[H]
%\begin{algorithmic}[1]
%\Procedure{Value Iteration Function}{$\beta, N, W, u(x), \varepsilon$, maxiter}
%    \State $w \gets [0, \frac{W}{N},\ldots, W]$
%      \Comment{Divide $W$ into $N$ equally sized pieces}
%    \State $V_0 \gets [V_0(w_0),V_0(w_1),\ldots,V_0(w_N)] $
%     \Comment{Common choice is $V_0(w_i)=u(w_i)$}
%    \For{$i=1,2,\dots,$\ \li{maxiter}}
%        \Comment Iterate only \li{maxiter} times at most.
%        \For{$w_i \in w$}
%            \State $V_{k+1}(w_i) = \max_{w_j \in \mathbf{w} \, | \, w_j \leq w_i} u(w_i-w_j) + \beta V_k(w_j)$
%        \EndFor
%        \If{ $||V_{k+1} - V_k|| < \varepsilon$}
%            \State \texttt{break}
%             \Comment{Stop iterating if the approximation stops changing enough.}
%        \EndIf
%    \EndFor
%    \State \pseudoli{return} $V_k$
%\EndProcedure
%\end{algorithmic}
%\caption{Value Function Iteration}
%\label{alg:ValueIteration}
%\end{algorithm}
%
%Most iterative algorithms have a \li{max\_iter} parameter that will terminate the algorithm after \li{max\_iter} iterations regardless of whether or not it has converged.
%This is because even though we have guaranteed convergence, we might have a convergence rate that is too slow to be useful.
%However, generally this algorithm will converge much faster than computing the true value function using dynamic programming as in Finite Value Iteration.
%
%\begin{warn} % Notation warning
%Note that $w_j \leq w_i$, because we can never have more cake at a later period; cake can only be consumed, not created.
%\end{warn}
%
%\begin{problem}
%\label{prob:policyiter-value1}
%Write a function called \li{value_iteration()} that will accept a numpy array representing the initial vector $V_0$, a discount factor $\beta \in (0,1)$,
%the number of states to split $\mathbf{w}$ into $N$, the amount of initial cake $W$, a utility function $u(x)$,
%the tolerance amount $\epsilon$, and the maximum number of iterations \li{max\_iter}.
%Perform value iteration until $\|V_{k+1} - V_{k}\| < \epsilon$ or $k > $ \li{max\_iter}.
%Return the final vector representing $V^*$.
%
%It is useful for our function to accept $V_0$ as a parameter instead of calculating an initial guess inside our function, so that we can try different initial states for $V_0$.
%
%Test your function with the parameters $N=400$, $\beta=.995$, $u(x)=\sqrt{x}$, $W=1$.
%Try different values for $V_0$ and see if you get the same value for $V^*(W)$ (The value should be approximately 9.4988).
%% Question to consider:
%How do different initial guesses for $V_0$ affect the number of iterations required for convergence?
%\end{problem}
%
%
%\subsection*{Policy Vector}
%
%The value function $V^*(w_i)$ that we found describes how much utility $w_i$ will yield over time, thus $V^*(W)$ is the optimal value for our problem:
%\begin{align*}
%  V^*(W) = \max_{c_t}  & \sum_{t=0}^T \beta^t u(c_t), \\
%  \mbox{subject to } & \sum_{t=0}^T c_t = W ,\\
%  & c_t \geq 0.
%\end{align*}
%Although $V^*(W)$ is the solution, it is usually more important to know which sequence of $(c_t)_{t=0}^T$ yields the solution.
%This sequence is known as the \emph{policy vector} $\mathbf{c} = [c_0, c_1, \ldots, c_T]$ that corresponds to eating $c_i$ cake at time $t=i$.
%
%If this were a truly infinite problem, $\mathbf{c}$ would be impossible to calculate; there would be an infinite number of time periods.
%Fortunately, in the cake-eating problem, we will never have more than $N$ time periods.
%This happens because it is never optimal to eat $0$ pieces of cake in a single time period.
%The discount factor $\beta$ means at least $1$ piece must be eaten at each step.
%It is important to note that the length of $\mathbf{c}$ changes depending on what $\beta$ is.
%When $\beta = 0$, for example, only the first time period matters because all the rest will yield zero utility.
%So when $\beta = 0$, $T = 0$ and the length of $\mathbf{c}$ is $1$.
%As the value of $\beta$ gets closer to $1$, $T$ increases as well.
%
%The policy vector, $\mathbf{c}$, is found by using the policy function: $\pi : \mathbb{R} \to \mathbb{R}$ which has the constraint $0 \leq \pi(w_i) \leq w_i \, \, \forall \, i.$
%$\pi(w_i)$ is the amount of cake to save for the next period, given we started with $w_i$ cake.
%In \eqref{eq:policyiter-val_iteration}, this was represented by $w_j$.
%
%Because $\pi(w_i)$ is the amount of cake saved until next period, we can use $V^*(W)$ and modify the Bellman equation to find $\pi$:
%\begin{equation}
%\label{eq:policyiter-pol_func}
%\pi(w_i) = \argmax_{w_j \in \mathbf{w} \, | \, w_j\leq w_i} u(w_i-w_j) + \beta V^*(w_j) \, \, \forall \, i.
%\end{equation}
%For our purposes, the policy function will be represented as a vector $\mathbf{\pi}$.
%This is convenient because in practice it is infeasible to code a functional representation for $\pi$.
%$\mathbf{\pi}_i$ dictates how many pieces of cake to save for the next period if we currently have $i$ pieces of cake.
%
%Once we have a vectorized representation of the policy function, $\pi$, we can use it to calculate the policy vector $\mathbf{c}$ using the relationship:
%\begin{align*}
%  c_t = w^{(t)} - \pi(w^{(t)})), \\
%  \mbox{where } w^{(t+1)} = \pi(w^{(t)}), \\
%  \mbox{and } w^{(0)} = W.
%\end{align*}
%
%For example, assume $w = [0,.1,.2,\ldots,1]$ and $\pi = [0, 0, 1, 2, 2, 3, 3, 4, 4, 5, 5]$.
%This means that if we have $7$ pieces of cake remaining, $\pi(7) = 4$ so we should save $4$ pieces of cake for the next time period.
%At $t=0$, $w^0=1=w_{10}$, so there are $10$ pieces of cake.
%$\pi(10) = 5$, so at time $0$, we want to save $5$ pieces of the cake.
%Then $c_0=10-5=5$, and we consume half of the cake.
%To calculate $c_1$, observe that $w^1= \pi(10)=5$.
%$\pi(w^1) = \pi(5) = 3$.
%So at time $2$, we eat $5-3=2$ pieces of cake, or $w_3=.2$ of the total cake.
%The resulting policy vector is $\mathbf{c} = [0.5,0.2,0.1,0.1,0.1]$.
%
%\begin{warn}
%$w^{(t)}$ represents a time index, how much cake we have at time $t$, not to be confused with $w_i$, the numeric value of having $i$ pieces.
%\end{warn}
%
%\begin{problem}
%Write a helper function called \li{extract_policy_vector()} that will accept an array of discretized values $\mathbf{w}$, and a vector representing a policy function $\pi$.
%Return the policy vector $\mathbf{c}$ that determines how much cake should be eaten at each time step.
%
%Test your function with $\mathbf{w} = [0,.1,.2,\ldots,1]$ and $\pi = [0, 0, 1, 2, 2, 3, 3, 4, 4, 5, 5]$, as in the example above.
%\end{problem}
%
%\begin{problem}
%Modify \li{value_iteration()} to return the true value function $V_{k+1}$ and the corresponding policy vector $\mathbf{c}$.
%\\(Hint: Use \eqref{eq:policyiter-pol_func} to find the policy function and then call
%\li{extract_policy_vector()} inside of \li{value_iteration()}).
%\end{problem}
%
%\section*{Policy Function Iteration}
%For infinite horizon dynamic programming problems, it can be shown that value function iteration converges relative to the discount factor, $\beta$.
%As $\beta$ approaches $1$, the number of iterations increases dramatically.
%As mentioned earlier $\beta$ is usually close to $1$, which means this algorithm can converge slowly.
%In Problem \ref{prob:policyiter-value1} you should have noticed that runtime was significantly longer to run for larger $N$ or
%$\beta$ closer to 1.
%
%In value iteration, we used an initial guess for the value function, $V_0$ and used (\ref{eq:policyiter-bellman2}) to iterate towards the true value function.
%Once we achieved a good enough approximation for $V^*$, we recovered the true policy function $\pi$.
%Instead of iterating on our value function, we can instead make an initial guess for the policy function, $\pi_0$, and use this to iterate toward the true policy function, called policy iteration.
%We do so by taking advantage of the definition of the value function, where we assume that our policy function yields the most optimal result.
%
%That is, given a specific policy function $\pi_k(W)$, we can modify \eqref{eq:policyiter-bellman2} by assuming that the policy function is the optimal choice.
%\begin{align*}
%  V_k(w_i) = \max_{w_j\in [0,w_i]} u(w_i-w_j) + \beta V_k(w_j) = u(w_i - \pi_k(w_i)) + \beta V_k(\pi_k(w_i)). \\
%\end{align*}
%This is because $w_j$ is the amount of cake we should have left over after $t=i$, which is $\pi(w_i)$.
%Because the value function is defined recursively,
%\begin{equation}
%  \label{eq:policyiter-val_from_policy}
%V_k(W) = \sum_{t=0}^{\infty}\beta^t u(\pi_k^t(W) - \pi_k^{t+1}(W)), \\
%\end{equation}
%where $\pi_k^t(W)$ means applying $\pi_k$ t times, and $\pi_k^0(W) = W$.
%Recall that \eqref{eq:policyiter-val_from_policy} will terminate in a finite number of steps (because we will eventually run out of cake to eat).
%Fortunately, in cake-eating, \eqref{eq:policyiter-val_from_policy} can alternatively be calculated by using dynamic programming, \eqref{eq:policyiter-bellman2} which defines the relationship as:
%\begin{equation}
%  \label{eq:policyiter-fast_val_from_policy}
%  V_k(W) = u(\pi_k^t(W) - \pi_k^{t+1}(W)) + \beta V_k(\pi_k^{t+1}(W)).
%\end{equation}
%
%This happens because $\pi_k^{t+1}(W) < W$, with the initial condition that $V_k(0) = 0$.
%Thus, given an initial guess for our policy function, $\pi_0$, we calculate the corresponding value function using (\ref{eq:policyiter-val_from_policy}), and then use \eqref{eq:policyiter-pol_func} to improve our policy function.
%The algorithm for policy function iteration can be summarized as follows:
%
%
%
%\begin{algorithm}[H]
%\begin{algorithmic}[1]
%\Procedure{Policy Iteration Function}{$\beta, N, W, u(x), \varepsilon$, maxiter}
%    \State $w \gets [0, \frac{W}{N},\ldots, W]$
%      \Comment{Divide $W$ into $N$ equally sized pieces}
%    \State $\pi_0 \gets [\pi_0(w_0),\pi_0(w_1),\ldots,\pi_0(w_N)] $
%     \Comment{Common choice is $\pi_0(w_i)=w_{i-1}$ with $\pi_0(0)=0$}
%    \For{$i=1,2,\dots,$\ \li{maxiter}}
%        \Comment Iterate only \li{maxiter} times at most.
%        \For{$w_i \in \mathbf{w}$}
%            \State $V_{k}(w_i) = u(\pi_k^t(w_i) - \pi_k^{t+1}(w_i)) + \beta V_k(\pi_k^{t+1}(w_i))$
%            \State $\pi_{k+1}(w_i) = \argmax_{w_j \in \mathbf{w} \, | \, w_j \leq w_i} u(w_i-w_j) + \beta V_k(w_j)$
%        \EndFor
%        \If{ $||\pi_{k+1} - V_k|| < \varepsilon$}
%            \State \texttt{break}
%             \Comment{Stop iterating if the approximation stops changing enough.}
%        \EndIf
%    \EndFor
%    \State \pseudoli{return} $V_k, \pi$
%\EndProcedure
%\end{algorithmic}
%\caption{Policy Iteration}
%\label{alg:PolicyIteration}
%\end{algorithm}
%

% TODO: This may be a more efficient way of finding the value function $V_k$. Should we implement this?
% In order to compute the value function, $V_k$ corresponding to a given policy $\pi_k$, we must solve
% \begin{equation}
% \label{Val_Fun}
% V_k(W) = u(W-W') + \beta V_k(W')
% \end{equation}
% for $V_k$.
%
% As always, we take a discrete approximation to $W$, obtaining a length $N$ vector
% \[
% w = (w_0, w_1, w_2, \ldots, w_N)
% \]
% giving the possible cake quantities.
% The variable $W'$ becomes $w' = \pi_k(w)$.
% Once we have done this, equation \eqref{Val_Fun} is a linear system which we can rewrite as
% \begin{equation*}
% V_k(w) = u(w-w') + \beta QV_k(w)
% \end{equation*}
% where $Q$ is the $N\times N$ matrix
% \begin{equation*}
% Q_{ij} = \left\{
%      \begin{array}{ll}
%        1 & \text{if} \quad  w_i' = w_j\\
%        0 & \text{otherwise.}
%      \end{array}
%    \right.
% \end{equation*}
% Solving this system of equations, we have $V_k(w) = (I-\beta Q)^{-1}u(w-w')$.
% Although $Q$ may be large, we can take advantage of the fact that it is sparse, containing only $N$ nonzero entries out of $N^2$ total entries.

%\begin{problem}
%\label{prob:cake_eating_policyfun}
%Write a function called \li{policy_iteration()} that will accept a numpy array representing the initial vector $\pi_0$, a discount factor $\beta \in (0,1)$,
%the number of states to split $\mathbf{w}$ into $N$, the initial amount of initial cake $W_{max}$, a utility function $u(x)$, the convergence tolerance $\epsilon$, and the maximum number of iterations \li{max_iter}.
%Perform Policy Iteration until $\|\pi_{k+1} - \pi_{k}\|_\infty < \epsilon$ or $n > $ \li{max_iter}.
%Return the final vector representing $V_k$ as well as the policy vector $\mathbf{c}$.
%
%Test your policy iteration by calling \li{value_iteration()} and \li{policy_iteration()} using the same values for $\beta$, $N$, $W_{max}$, and $u(x)$.
%The value functions returned should be close to equal (use \li{np.allclose()}), and the policy vectors $\mathbf{c}$ should be identical.

% How many iterations does Policy Iteration take to converge?
% How does this compare to the number of iterations Value Iteration takes?

% TODO: Is this a more efficient way to do it?
% You will still need to pre-compute all values of $u(W-W')$, storing these in an $N \times N$ array. Make sure, as before,
% that the upper triangular entries are set to large negative values, so that we don't choose to consume more cake than is
% available. In the code snippets below, we will refer to this array as \li{U}.
%
% You may also find it convenient to not track the approximated policy function $\psi_k$ directly during the iteration, but
% rather to have a length $N$ array \li{psi_ind}, whose $i$-th entry gives the index of $w_i'$ relative to the discrete approximation
% $w$. Thus, rather than initializing a policy function $\psi_0$ directly, you can initialize it indirectly by setting, for example,
% \begin{lstlisting}
% >>> psi_ind = np.arange(N)
% \end{lstlisting}
% which corresponds to an initial policy $\psi_0(W) = W$.
% The policy function vector can be obtained from \li{psi_ind} by simply slicing \li{w} as follows:
% \begin{lstlisting}
% >>> psi = w[psi_ind]
% \end{lstlisting}
%
% In order to take advantage of the sparse matrices $I$ and $Q$, use the following imports from the SciPy \li{sparse} library.
% \begin{lstlisting}
% >>> from scipy import sparse
% >>> from scipy.sparse import linalg
% \end{lstlisting}
% and the following code to initialize $I$ (outside the loop)
% \begin{lstlisting}
% >>> I = sparse.identity(N, format='csr')
% >>> rows = np.arange(0, N)
% \end{lstlisting}
% and $Q$ (inside the loop)
% \begin{lstlisting}
% >>> columns = psi_ind
% >>> data = np.ones(N)
% >>> Q = sparse.coo_matrix((data, (rows,columns)), shape=(N,N))
% >>> Q = Q.tocsr()
% \end{lstlisting}
% Rather than compute $(I-\beta Q)^{-1}$ directly, use Scipy's sparse solver
% \begin{lstlisting}
% V = linalg.spsolve(I-beta*Q, u(w-w[psi_ind]))
% \end{lstlisting}
% where \li{u} is the square root function.
%
% In each iteration, we update \li{psi_indices} much as we did in the value function iteration algorithm:
% \begin{lstlisting}
% >>> psi_ind = np.argmax(U + beta*V, axis=1)
% \end{lstlisting}
% where we assume \li{V} has shape \li{(1,N)}.
%
% Use the 2-norm when computing $\|\psi_{k+1} - \psi_k\|$, and set the tolerance to $\delta = 1e-9$.
%
% Take $N = 1000$ and $\beta = .95$.
%\end{problem}

%%%%%%%%%%%%%% MPI %%%%%%%%%%%%%%%%%%%% Uncomment when we fix it.
% \section*{Modified Policy Function Iteration}
% Although policy iteration converges in fewer iterations, computing $V_k$ using (\ref{eq:policyiter-val_from_policy}) at each iteration can be quite costly, especially for problems with large $N$.
% Thus policy iteration converges slowly relative primarily to $N$, and value iteration converges slowly relative primarily to $\beta$.
% It turns out there is a hybrid method called modified policy function iteration that seeks to simultaneously minimize these disadvantages by performing policy iteration and value iteration together.
%
% Modified policy iteration is identical to policy iteration, with the addition of one extra step.
% Before \ref{alg:computing_pi_k} of the policy iteration algorithm, take the $V_k$ computed in step \ref{alg:computing_V_k} and perform value iteration with \li{max_iter}$=m$ using $V_k$ as the initial guess.
% This is faster than solving for the exact value function for large state spaces at every single iteration.
% There is no strict rule on the value of $m$, but typically $m \in [5,15]$ work well.
%
% Thus the algorithm for modified policy iteration can be summarized as follows:
% \begin{enumerate}
%   \item Discretize the space into $N$ equal sized pieces: $\mathbf{w} = \left[w_0, w_1, \ldots, w_N\right]$ where $w_0 = 0, \, \, w_N = W_{max}$.
%   \item Choose an initial vector to represent $\pi_0$: $\left[\pi_0(w_0), \pi_0(w_1), \ldots, \pi_0(w_N)\right]$
%             A common initial choice for $\pi_0$ is $\pi_0(w_i) = w_{i-1}$, meaning we save $w_{i-1}$ pieces, and $\pi_0(w_0) = 0$.
%   \item For $k = 1, 2, \ldots,$ \li{max\_iter}:
%   \begin{enumerate}
%     \item \label{alg:computing_V_k2} For each $w_i \in \mathbf{w}$ compute the value function $V_k(w_i)$ using (\ref{eq:policyiter-val_from_policy}).
%     \item \label{alg:value_iteration2} Call \li{value_iteration()} with $V_k$ as the initial guess, and use $max\_iter = m$, for some $m$.
%         Use the value function that it returns ($\hat{V_k}$) as the value function to use in \ref{alg:computing_pi_k2}.
%     \item \label{alg:computing_pi_k2} For each $w_i \in \mathbf{w}$ calculate $\pi_{k+1}(w_i) = \argmax_{W' \in \mathbf{w} \, | \, W' \leq w_i} u(w_i-W') + \beta \hat{V_k}(W')$.
%     \item If $\| \pi_{k+1} - \pi_k \| < \epsilon,$ terminate early.
%   \end{enumerate}
% \end{enumerate}
% Note that our previous methods, value iteration, and policy iteration are like two sides of our algorithm: iterating on the value function or iterating on the policy function.
% Modified policy function does a combination of both, attempting to take advantage of both methods while trying to avoid their disadvantages.
% Because modified policy iteration takes only slightly more work to code compared to value iteration and policy iteration, it is often preferred in practice.
% Modified policy iteration usually performs better than value iteration or policy iteration.
%
% \begin{problem}
% Write a function \li{modified_policy_iteration()} that takes the same arguments as \li{policy_iteration()}, along
% with an additional keyword argument $m$, which indicates the number of iterations of value iteration to do at each iteration of policy iteration.
% Return the converged value function and policy vector.
%
% Much of the code in \li{policy_iteration()} can be re-used, you just need to implement the extra step introduced by modified policy iteration.
% This is why the \li{value_iteration()} function accepts $V_0$ as a parameter, it makes implementing MPI much simpler.

% TODO: Based on other TODOS, do we want to solve the linear system?
% The key differences are as follows.
% First, you do not need to initialize the sparse arrays $I$ and $Q$ in the modified policy function iteration.
% Secondly, rather than solving a linear system of equations to obtain the value function, you will loop
% the equation
% \begin{lstlisting}
% >>> V = u(w - w[psi_ind]) + beta*V[psi_ind]
% \end{lstlisting}
% $m$ times in each iteration. Notice how this line of code corresponds with Equation \eqref{Val_Fun}, where
% \li{w} represents $W$, \li{w[psi_ind]} represents $W'$, and \li{V[psi_ind]} represents $V(W')$.
%
% The remainder of the code should be unchanged.
%\end{problem}
%%%%%%%%%%%%%%%%%%%% END OF MPI SECTION AND PROBLEM %%%%%%%%%%%%%%%%%%%%%%

%\begin{problem}
%Solve the cake eating problem with both value iteration and policy iteration for various values of $\beta$ and compare how long each method takes.
%% Each final $V^*$ should be \emph{np.allclose} and your policy vectors $\mathbf{c}$ should be identical for both methods.
%Use $N=1000$ as the number of grid points for $\mathbf{w}$ and $\beta = [.95, .975, .99, .995]$.
%
%It is important to use feasible initial guesses in each case in order to make the results comparable.
%A good initial guess greatly affects the number of iterations required for convergence.
%Use $V_0(w_i) = u(w_i)$ for value iteration, and $\pi_0(w_i) = w_{i-1}$, with $\pi_0(w_0) = 0$ for policy iteration.
%\\(Hint: set \li{max\_iter} high enough for each method so that the functions actually converge; large values of $\beta$ may require several hundred iterations for value iteration.)
%
%Graph the results for each method with $\beta$ on the x-axis and time on the y-axis.
%Compare your results to the following figure.
%
%\begin{figure}[H] % This figure plots the time each method takes to converge
%\centering        % use width=.7\textwidth to size the figure.
%    \includegraphics[width=.7\textwidth]{figures/comparing_methods.pdf}
%    % \caption{Comparing each method for various values of $\beta$.}
%    \label{fig:basic1}
%\end{figure}
%
%% You were not asked to compare these functions to the \li{find_policy()} function written in Finite Value Iteration because it is significantly slower than these iterative methods.
%\end{problem}

\begin{comment}
\newpage

\section*{Additional Material}
Perhaps here we can talk a bit about truly infinite MDPs such as choosing an optimal path through some kind of grid.

Or we could spend time talking about Stochastic MDPs where we introduce probabilities into our decision making process.

Or we could try and answer a question like this:
Given the following grid, we want to know if it's going to be more valuable to stop and get the flags or if it'll be better to head straight to the goal.
The answer is dependent on what kind of reward we get for touching the flags, the reward for reaching the goal, and the discount factor.
\begin{figure}[H] % Describe the figure here.
\centering        % use width=.7\textwidth to size the figure.
    \includegraphics[width=.7\textwidth]{figures/flag_maze.pdf}
    \caption{Source: \url{http://www.samyzaf.com/ML/rl/qmaze.html}}
    \label{fig:basic1}
\end{figure}

\end{comment}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Remnants from a Previous Version of the lab
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% \section*{Discrete Choice (Threshold) Problems}\label{SecDiscrChoice}
%
% One powerful application of dynamic programming is that we can make
% models that have both continuous and discrete state variables. These models are sometimes referred to as
% discrete choice problems or optimal stopping problems.  Examples include models of employment that involve both
% the choice of whether to work and how much to work, models of firm entry and exit that involve the choice of both
% whether to produce and how much to produce, and models of marriage that involve the choice of whether to date
% (get married or keep dating) and how much to date.
% This application illustrates the versatility of dynamic programming as a dynamic solution method
%
% In this lab, we follow a simple version of a standard job search model.
% Assume that workers live infinitely long.
% We will split a worker's life into discrete time periods, and in each period the worker is either
% employed or unemployed, and receives a job offer. The worker must make a choice between discrete actions
% (such as accepting or rejecting a job offer), with the goal of maximizing some utility function (hence,
% this is a \emph{discrete choice} problem).
%
% We can state this problem in terms of dynamic programming by defining an appropriate value function.
% Let the value of entering a period with most recent wage $w$,
% current job offer wage $w'$, and employment status $s$ be given by the following value function,
% \begin{equation}\label{EqV}
%    V(w,w',s) = \begin{cases}
%                   V^E(w)    \quad&\text{if}\quad s = E \\
%                   V^U(w,w') \quad&\text{if}\quad s = U \\
%                \end{cases}
% \end{equation}
% where employment status is a binary variable $s\in\{E,U\}$ ($E$ indicates ``employed" and $U$ indicates ``unemployed");
% a person can be either employed or unemployed.
%
% As in the cake eating problem, the value function is calculated as the sum of some reward (based on the
% current state) and the discounted value of entering the next period in some particular state.
% The reward function, denoted (as usual) by $u$, gives the utility of spending available funds.
% Assuming that a worker receives some wage $x$ in a given period and spends all available money in the period,
% the utility of consumption is given by
% \[
% u(x).
% \]
% Calculating the value for the next period depends on the employment status in the current period, so we address
% this separately for each case.
%
% Let us first consider the case where the individual is unemployed ($s = U$). As is customary, let $s'$ denote the
% employment status of the worker in the next period.
% In this unemployed state, the worker receives unemployment benefits equal to a fraction of her most recent wage,
% i.e. $\alpha w$, where $\alpha \in (0, 1)$.
% Hence, the utility of consumption in the current unemployed state is given by
% \[
% u(\alpha w).
% \]
%
% The worker also receives one wage offer ($w'$) per period, and will obtain
% this wage in the next period provided that she chooses to accept employment, i.e. provided $s' = E$.
% The worker must decide whether to accept the current wage offer $w'$ or to remain unemployed in the next period,
% i.e. she must decide on the value of $s'$. How
% does she make this choice? She must weigh the value of entering the next period as an employed worker with wage
% $w'$ (given by $V^E(w')$) versus the value of entering the next period as an unemployed worker with
% previous wage $w$ and unknown wage offer $w''$ (given by $V^U(w,w'')$). Because the worker cannot know
% what the future wage offer $w''$ will be, it is treated as a random variable with a particular probability
% distribution. Hence, the worker must actually compute the \emph{expected} value of entering the next
% period unemployed. This term is simply
% \[
% \mathbb{E}_{w''}V^U(w,w''),
% \]
% where $\mathbb{E}_{w''}$ denotes the expectation operator with respect to the probability distribution of future
% wage offers $w''$.
% To sum up, the worker chooses to accept the wage offer $w'$ or remain unemployed in the next period based on
% which option gives the greater expected value, and the value of this decision is given by
% \[
% \max\Bigl\{V^E(w'), \,\, \mathbb{E}_{w''}V^U(w,w'')\Bigr\}.
% \]
%
% The overall value of the current unemployed state with previous wage $w$ and current wage offer $w'$ is
% just the utility of consumption plus the discounted value of the next period, i.e.
% \begin{equation}\label{EqVu}
% V^U(w,w') = u(\alpha w) + \beta \max\Bigl\{V^E(w'), \,\, \mathbb{E}_{w''}V^U(w,w'')\Bigr\},
% \end{equation}
% where $\beta$ is the discount factor.
%
% Now we turn to the case where the job status is employed ($s = E$).
% In this case, the worker receives a wage $w$ in the current period, and so the utility of consumption is
% just
% \[
% u(w).
% \]
% In the next period, the worker will have most recent wage $w$, she will receive wage offer $w''$, and will
% have employment status $s'$. As in the unemployed case, $w''$ is unknown and treated as a random variable.
% Unlike the unemployed case, however, the worker's future employment status $s'$ is not under her control,
% but rather is also a random variable. The reason for this is that the worker will remain employed
% until she loses the job, a random event that occurs with some fixed probability in each time period.
% Hence, we must calculate the expected value of the next period with respect to both $w''$ and $s'$.
% We may write the entire value function for the employed case as
% \begin{equation}\label{EqVe1}
%    V^E(w) = u(w) + \beta \mathbb{E}_{w'',s'}V(w,w'',s').
% \end{equation}
%
% To calculate the expectation term, we need to know the joint probability distribution over $w''$ and $s'$.
% This can be characterized in the following way.
% We assume that $s'$ and $w''$ are independent. Hence, we can split the joint expectation
% operator into the composition of the two individual expectation operators:
% \[
% \mathbb{E}_{w'',s'} = \mathbb{E}_{w''}\mathbb{E}_{s'}.
% \]
% Let $\gamma$ represent the probability that an employed worker becomes unemployed in the next period,
% so that $1-\gamma$ is the probability of remaining employed in the next period.
% If the worker stays employed in the next period ($s' = E$), then next period's wage equals the current
% period's wage, and the term inside the expectation is
% \[
% V(w,w'',E) = V^E(w).
% \]
% We then have
% \begin{align*}
% \mathbb{E}_{s'}V(w,w'',s') &= (1-\gamma)V(w,w'',E) + \gamma V(w,w'',U)\\
% &= (1-\gamma)V^E(w) + \gamma V^U(w,w'').
% \end{align*}
% Notice that the term $(1-\gamma)V^E(w)$ is constant with respect to $w''$. Then
% \begin{align*}
% \mathbb{E}_{w''}\mathbb{E}_{s'}V(w,w'',s') &= \mathbb{E}_{w''}\left[(1-\gamma)V^E(w) + \gamma V^U(w,w'')\right]\\
% &= \mathbb{E}_{w''}(1-\gamma)V^E(w) + \mathbb{E}_{w''}\gamma V^U(w,w'')\\
% &= (1-\gamma)V^E(w) + \gamma \mathbb{E}_{w''}V^U(w,w'').
% \end{align*}
% Hence, we can rewrite \eqref{EqVe1} as follows:
% \begin{equation}\label{EqVe2}
%    V^E(w) = u(w) + \beta \Bigl[(1-\gamma)V^E(w) + \gamma \mathbb{E}_{w''}V^U(w,w'')\Bigr].
% \end{equation}
%
% We have now completely described the value function. What about the policy function?
% The policy function for the unemployed worker gives her decision on whether to accept the job $s'=E$
% or to reject the job $s'= U$.
% This will be a function of both the most recent wage $w$ and the current wage offer $w'$.
% The employment status $s'$ in the next period is determined by the policy function $\psi$:
% \[
% s' = \psi(w,w').
% \]
%
% These discrete choice problems are often called threshold
% problems because the policy choice depends on whether the state variable is greater than or less than
% some threshold level. That is, an unemployed worker will accept a job if and only if the offer wage is
% above some set amount that depends on the most recent wage $w$. In the labor search model,
% the threshold level is called the ``reservation wage'' $w_R'$. The reservation wage $w_R'$ is defined as
% the wage offer such that the worker is indifferent between accepting the job $s' = E$ and
% staying unemployed $s' = U$. Hence, this reservation wage satisfies the equation
% \begin{equation}\label{EqWR}
%    V^E(w_R') = E_{w''}\left[V^U(w,w'')\right].
% \end{equation}
% The policy function will then take the form of accepting the job if $w' \geq w_R'$ or
% rejecting the job offer and remaining unemployed if $w' < w_R'$:
% \begin{equation}\label{EqSprime}
%    s' = \psi(w,w') = \begin{cases}
%                       E \quad\text{if}\quad w' \geq w_R' \\
%                       U \quad\text{if}\quad w' < w_R'.
%                    \end{cases}
% \end{equation}
% Figure \ref{fig:disc_policy} shows an example of the discrete policy function.
%
% \begin{figure}
% \includegraphics[width=\textwidth]{figures/disc_policy.pdf}
% \caption{Here is the policy function for fixed $w = 50$.  Numerically we let 0 represent unemployment, $U$,
% and 1 represent employment, $E$.  Thus we see that an individual will choose to take a new job, given
% their old wage was 50, at a wage of roughly 35.  Thus for a previous wage of 50, we say the reservation wage is 35.}
% \label{fig:disc_policy}
% \end{figure}
%
% In summary, the labor search discrete choice problem is characterized by the value functions \eqref{EqV}, \eqref{EqVu},
% and \eqref{EqVe2}, the reservation wage \eqref{EqWR}, and the policy function \eqref{EqSprime}. Because wage offers
% are distributed according to some given probability distribution (denote the cdf by $F(w')$),
% and because the policy function takes the form of \eqref{EqSprime},
% the probability that the unemployed worker receives a wage offer that she will reject is $F(w_R')$ and the probability
% that she receives a wage offer that she will accept is $1 - F(w_R')$. Just like the continuous-choice cake eating
% problems, this problem can be solved by value function, policy function, or modified policy function iteration.
%
% The value function iteration solution method for the equilibrium in the labor search problem is analogous to the
% value function iteration from the previous labs. The only difference is that two value functions ($V^E$ and $V^U$)
% must converge to a fixed point in this problem instead of just one value function converging in the previous problems.
% Although there are two value functions to consider, there is only one policy function, since decisions are only made
% in the unemployed state.  Thus, there is only one policy function on which to iterate in the case of policy or modified
% policy iteration.
%
% In the following problems, you will solve the job search problem using value function iteration and modified policy
% function iteration. Assume that the consumption utility function $u$ is given by
% \[
% u(w) = \sqrt{w}.
% \]
% Assume that the probability of becoming unemployed in a given period is $\gamma = 0.10$, the fraction of wages paid
% in unemployment benefits is $\alpha = 0.5$, and the discount factor is $\beta = 0.9$.
%
% Assume that the log of wage offers are distributed normally.  We then say that offers are distributed
% lognormally and write
% \[
% w'\sim \text{LogN}(\mu,\sigma).
% \]
% This is a convenient choice for the distribution of
% wage offers.  Among other things, it guarantees that wage offers will be positive.
% A mean of $20$ and variance of $200$ are typical parameters of such a wage distribution,
% and we will use these parameters in the following problems.
%
% As usual when dealing with continuous variables, we form a discrete approximation of
% the possible wage values. In particular, approximate the wage values by a vector of
% length $N = 500$ of equally-spaced values from $w_{min} = 0$ to $w_{max} = 100$, inclusive.
% We then form a corresponding discrete approximation of the probability density function
% $f(w')$ for the lognormal wage offers using code provided in the file \li{discretelognorm.py},
% as follows, where \li{w} is the length-$N$ vector of wage values, $m$ is the mean, and $v$ is
% the variance as specified above:
% \begin{lstlisting}
% >>> from discretelognorm import discretelognorm
% >>> f = discretelognorm(w, m, v)
% \end{lstlisting}
% The function \li{discretelognorm} computes the discrete pdf of the specified lognormal distribution in much
% the same way that you calculate the discrete normal pdf when solving the stochastic cake eating problem.
%
%
% \begin{problem}
% Solve the job search problem using value function iteration. Return the converged value functions
% $V^E$ and $V^U$, as well as the converged policy function $\psi$. The following steps provide detailed
% instructions. Note that there are multiple ways to proceed, and the following is simply one (fairly good)
% possibility.
% \begin{enumerate}
%
%    \item As described above, represent the possible wage values by an array \li{w} of length
%    $N$. Denote this array entrywise by
%    \[
%    w = (w_1,w_2,\ldots,w_N).
%    \]
%    Calculate the corresponding discrete lognormal pdf \li{f}, exactly as shown above.
%
%    \item Note that $u(w)$ and $u(\alpha w)$ are needed when computing the value functions.
%    Since these quantities do not change from one iteration to another, it is smart to
%    compute them once at the outset. Denote $u(w)$ by \li{uw} and $u(\alpha w)$ by \li{uaw}.
%    These are easily calculated as follows:
% \begin{lstlisting}
% >>> uw = u(w)
% >>> uaw = u(alpha*w).reshape((N,1))
% \end{lstlisting}
%     where \li{u} is the square root function. We must reshape \li{uaw} because of array broadcasting
%     issues that arise in the code snippets below.
%
%    \item Since $V^E$ is a function of only $w$, it will be represented by a vector of length $N$, where
%    the $i$-th entry gives $V^E(w_i)$. The unemployed value function $V^U$, however,
%    is a function of both $w$ and $w'$, so it will be represented by an $N \times N$ array,
%    where the $(i,j)$-th entry gives $V^U(w_i,w_j)$. Initialize the entries of these arrays to 0:
% \begin{lstlisting}
% >>> VE = np.zeros(N)        #employed value function
% >>> VU = np.zeros((N,N))    #unemployed value function
% \end{lstlisting}
%
%    \item Note that $\mathbb{E}_{w''}V^U(w,w'')$  is needed to calculate both $V^E$ and $V^U$.
%    This expectation depends on $w$, and so can be represented by a length $N$ array, where the
%    $i$-th entry is $\mathbb{E}_{w''}V^U(w_i,w'')$.
%    It is convenient to assign a variable to this array to keep track of it throughout the iterations.
%    We denote the expectation by \li{EVU}, and initialize it to zeros:
% \begin{lstlisting}
% >>> EVU = np.zeros(N)
% \end{lstlisting}
%
%    \item For reasons that will soon become apparent, we will need to create an $N\times N$ helper array
%    whose rows are equal to \li{VE} (call this array \li{MVE}), and a $N \times N$ helper array whose columns
%    are equal to \li{EVU} (call this array \li{MEVU}).
%    At the outset, simply initialize these arrays to zeros.
%
%    \item Because job status is a binary variable, the policy function returns one of two possible values. It is
%    convenient to represent ``employed" by $1$ and ``unemployed" by $0$. Now the policy function depends
%    on $w$ and $w'$, so it will also be represented by an $N\times N$ array \li{PSI} of zeros and ones,
%    where the $(i,j)$-th entry gives $\psi(w_i, w_j)$.
%
%    \item Now we are ready to begin the iteration.
%    A single iteration involves computing the updated value functions $V^E$ and $V^U$ from
%    equations \eqref{EqVe2} and \eqref{EqVu} and then calculating the $2$-norm distance between
%    both pairs of old and updated value functions to test for convergence. If both of these
%    $2$-norm distances are less than $10^{-9}$, terminate the iteration.
%
%    Before calculating the updated value functions, we first update our helper arrays \li{MVE} and
%    \li{MEVU}. The rows of \li{MVE} need to equal \li{VE}. We can use array broadcasting:
% \begin{lstlisting}
% >>> MVE[:,:] = VE.reshape((1,N))
% \end{lstlisting}
%    The columns of \li{MEVU} need to equal \li{EVU}, so use a similar technique:
% \begin{lstlisting}
% >>> MEVU[:,:] = EVU.reshape((N,1))
% \end{lstlisting}
%
%    Now let us address how to compute the updated $V^U$, which we denote by \li{VU1}.
%    Equation \eqref{EqVu} shows that it is the sum of
%    two terms. The first, $u(\alpha w)$, we have already computed and stored in the variable \li{uaw}.
%    The second term involves a maximization between two alternatives. One can imagine writing
%    a double for loop ranging over the values of $w'$ and $w$ to compute each individual
%    $\max\{V^E(w'), \mathbb{E}_{w''}V^U(w,w'')\}$, but we can take advantage of the helper
%    arrays \li{MVE} and \li{MEVU} to do this computation in one efficient line of code.
%    Note that the $(i,j)$-th entry of \li{MVE} is just $V^E(w_j)$ and the $(i,j)$-th
%    entry of \li{MEVU} is $\mathbb{E}_{w''}(w_i, w'')$, and
%    \[
%    V^U(w_i,w_j) = u(\alpha w_i) + \beta\max\{V^E(w_j), \mathbb{E}_{w''}V^U(w_i,w'')\}.
%    \]
%    Hence, taking the entrywise maximum of the arrays \li{MVE} and \li{MEVU} gives us the appropriate
%    max term for $V^U$. To get the entrywise maximum of two arrays, stack the arrays along a new
%    axis using \li{np.dstack}, and maximize along that axis. The computation for \li{VU}, then, is
%    \begin{lstlisting}
% >>> VU1 = uaw + beta*np.max(np.dstack([MEVU, MVE]), axis=2)
%    \end{lstlisting}
%
%    Calculating the updated $V^E$, denoted by \li{VE1}, is more straightforward.
%    Equation \eqref{EqVe2} shows that it is just
%    a particular linear combination of the arrays \li{uw}, \li{VE}, and \li{EVU}:
%    \begin{lstlisting}
% >>> VE1 = uw + beta*((1-gamma)*VE + gamma*EVU)
%    \end{lstlisting}
%
%    We can now calculate the 2-norm distances between old and updated value functions.
%    It remains to update \li{VE}, \li{VU}, and \li{EVU}.
%    The first two updates are trivial, and calculating \li{EVU} is equivalent to the
%    matrix-vector multiplication of \li{VU} with \li{f}. This is similar to how we computed
%    expectations in previous labs:
%    \begin{lstlisting}
% >>> EVU = np.dot(VU,f).ravel()
%    \end{lstlisting}
%    We use the \li{ravel} function to ensure that \li{EVU} is a flat array.
%
%    \item Notice that it is not necessary to iteratively update the policy function, as it is not needed
%    to update the value functions. Thus, we need only compute the policy function once, after convergence
%    of the value functions has been achieved. This is done in a manner similar to calculating
%    $\max\{V^E(w'), \mathbb{E}_{w''}V^U(w,w'')\}$ as described above, except we need to take the \emph{argmax}:
% \begin{lstlisting}
% >>> PSI = np.argmax(np.dstack([MEVU,MVE]), axis=2)
% \end{lstlisting}
%
%    \item Compute the reservation wage $w_R'$ as a function of the current wage $w$. It will be represented
%    by a length $N$ array called \li{wr}. The reservation wage is the
%    value of $w'$ where the policy function changes from zeros to ones (the optimal choice changes from remaining
%    unemployed to accepting the job offer). We can calculate this as follows:
%    \begin{lstlisting}
% >>> wr_ind = np.argmax(np.diff(PSI), axis = 1)
% >>> wr = w[wr_ind]
%    \end{lstlisting}
%
%    \item Plot the equilibrium reservation wage $w_R'$ of the converged problem as a function of the current
%    wage $w$ with the current wage on the $x$-axis and the reservation wage $w_R'$ on the $y$-axis. This is
%    the most common way to plot discrete choice policy functions. The reservation wage represents the wage
%    that makes the unemployed worker indifferent between taking a job offer and rejecting it. So any wage
%    above the reservation wage line represents $s' = E$ and any wage below the reservation wage line represents
%    $s' = U$. Your plot should resemble that in Figure \ref{fig:res_wage}.
%
% \end{enumerate}
% \end{problem}
%
% \begin{figure}
% \includegraphics[width=\textwidth]{figures/reservation_wage.pdf}
% \caption{The reservation wage as a function of previous wage for the job search problem.}
% \label{fig:res_wage}
% \end{figure}
%
% In the previous problem, it was necessary to iterate on two value functions.
% Consequently the convergence is relatively slow.
% We can improve upon this situation by using modified policy function iteration.
%
% \begin{problem}
% Solve the same problem, this time using modified policy function iteration with $15$ value function iterations
% within each policy iteration.  You should be able to re-use much of your code from the previous problem.
%
% Start off by initializing all of the same variables. Additionally, initialize your policy function array \li{PSI},
% say
% \begin{lstlisting}
% >>> PSI = 2*np.ones((N,N))
% \end{lstlisting}
%
% Next comes the iteration.
% Essentially, the iteration will consist of an outer while-loop (which terminates once the 2-norm distance
% between successive policy functions passes below $10^{-9}$), and an inner for loop (with $15$ loops).
%
%
% The first step in the while-loop is to calculate the new policy function \li{PSI1}, just as in the previous
% problem. Next, perform the inner for-loop, which consists simply of the value function iteration, but this
% time using the current policy function. This means the line of code
% \begin{lstlisting}
% >>> VU = uaw + beta*np.max(np.dstack([MEVU, MVE]), axis=2)
% \end{lstlisting}
% is no longer valid, as it does not use the policy function. We must instead have
% \begin{lstlisting}
% >>> VU = uaw + beta*(MVE*PSI1 + MEVU*(1 - PSI1))
% \end{lstlisting}
% Why is this code correct?
%
% Finally, after exiting the for loop, calculate the 2-norm distance between the old and the new policy function,
% and then update your old policy function, i.e.
% \begin{lstlisting}
% >>> PSI = PSI1
% \end{lstlisting}
%
% After convergence is achieved, once again compute the reservation wage array, and plot it as in the previous
% problem. Then return the converged policy function.
% \end{problem}
%
% \begin{problem}
% How many iterations did the value function iteration method take?
% How many iterations did the modified policy function iteration method take?
% Which was faster?
% \end{problem}