\begin{center}
  \begin{tikzpicture}
    \scriptsize
    \begin{scope}
      \node [draw=violet,fill=violet!5,rounded corners,minimum width=6.5cm,minimum height=3cm] (b1) at (0,0) {};
      \node [anchor=north,text width=6.4cm,align=center] (b11) at ([yshift=-0.1cm]b1.north)
      {Teacher's next-token distribution given the state\\${\textcolor{blue}{x},\textcolor{orange}{y_{<i}}}$};
      \draw [anchor=north,violet] ([xshift=.2cm,yshift=.6cm]b1.south west) -- ([xshift=-.2cm,yshift=.6cm]b1.south east);
      \node [draw=violet,anchor=south,fill=violet!30,minimum width=.75cm, minimum height=.4cm] (b12) at ([xshift=.3cm,yshift=.6cm]b1.south) {};
      \node [draw=violet,anchor=south east,fill=violet!30,minimum width=.75cm, minimum height=1cm] (b13) at ([xshift=-.6cm]b12.south west) {};
      \node [draw=violet,anchor=south east,fill=violet!30,minimum width=.75cm, minimum height=.3cm] (b14) at ([xshift=-.6cm]b13.south west) {};
      \node [draw=violet,anchor=south west,fill=violet!30,minimum width=.75cm, minimum height=.1cm] (b15) at ([xshift=1.2cm]b12.south east) {};
      \node [anchor=south] (b16) at ([yshift=.1cm]b13.north) {\textcolor{violet}{$x=5$}};
      \node (b17) at (b16-|b12) {$x=6$};
      \node (b18) at (b16-|b15) {(eos)};
      \node (b19) at ($1/2*(b17)+1/2*(b18)$) {$\cdots$};
      \node (b20) at (b16-|b14) {$x=4$};
      \node [anchor=north] (btitle) at ($1/2*(b14.south)+1/2*(b15.south)$){$P^t(\cdot|x,y_{<i})$};
    \end{scope}

    \begin{scope}[yshift=-3.5cm]
      \node [draw=orange,fill=orange!5,rounded corners,minimum width=6.5cm,minimum height=3cm] (b2) at (0,0) {};
      \node [anchor=north,text width=6.4cm,align=center] (b21) at ([yshift=-0.1cm]b2.north)
      {Student's next-token distribution given \\ the same state\\${\textcolor{blue}{x},\textcolor{orange}{y_{<i}}}$};
      \draw [anchor=north,orange] ([xshift=.2cm,yshift=.6cm]b2.south west) -- ([xshift=-.2cm,yshift=.6cm]b2.south east);
      \node [draw=orange,anchor=south,fill=orange!30,minimum width=.75cm, minimum height=.4cm] (b22) at ([xshift=.3cm,yshift=.6cm]b2.south) {};
      \node [draw=orange,anchor=south east,fill=orange!30,minimum width=.75cm, minimum height=.5cm] (b23) at ([xshift=-.6cm]b22.south west) {};
      \node [draw=orange,anchor=south east,fill=orange!30,minimum width=.75cm, minimum height=.8cm] (b24) at ([xshift=-.6cm]b23.south west) {};
      \node [draw=orange,anchor=south west,fill=orange!30,minimum width=.75cm, minimum height=.1cm] (b25) at ([xshift=1.2cm]b22.south east) {};
      \node [anchor=south] (b26) at ([yshift=.1cm]b24.north) {$x=4$};
      \node (b27) at (b26-|b22) {$x=6$};
      \node (b28) at (b26-|b25) {(eos)};
      \node (b29) at ($1/2*(b27)+1/2*(b28)$) {$\cdots$};
      \node (b20) at (b26-|b23) {\textcolor{violet}{$x=5$}};
      \node [anchor=north] (btitle) at ($1/2*(b24.south)+1/2*(b25.south)$){$P^s_\theta(\cdot|\textcolor{orange}{x},\textcolor{orange}{y_{<i}})$};
    \end{scope}

    \draw[->,dashed] (b24|-b2.north) -- ([yshift=-0.1cm]b14.south);
    \draw[->,dashed] (b23|-b2.north) -- ([yshift=-0.1cm]b13.south);
    \draw[->,dashed] (b22|-b2.north) -- ([yshift=-0.5cm]b12.south);
    \draw[->,dashed] (b25|-b2.north) -- ([yshift=-0.1cm]b15.south);

    \node[draw,fill=violet!10,rounded corners,anchor=north west,align=center,minimum width=2cm,minimum height=1.5cm] (tm) at ([xshift=2cm]b1.north east) {\small Teacher \\[.1cm] \small Model};
    \node[align=center] (tmin) at ([yshift=-1cm]tm.south) {Input (state): \\ $\textcolor{blue}{x},\textcolor{orange}{y_{<i}}$};
    \node[draw,fill=orange!10,rounded corners,anchor=south east,align=center,minimum width=2cm,minimum height=1.5cm] (sm) at ([xshift=-2cm]b2.south west) {\small Student \\[.1cm] \small Model};
    \node[align=center] (smin) at ([yshift=-1cm]sm.south) {Input (problem): \\[.1cm] \small$\textcolor{blue}{x}$};
    \draw[->] (tmin) -- (tm); \draw[->] (tm) -- (tm-|b1.east);
    \draw[->] (smin) -- (sm); \draw[->] (sm) -- node[above]{Generate} (sm-|b2.west);

    \node[draw,anchor=north,rounded corners,minimum width=5.5cm, minimum height=1.5cm,dashed] (b3) at ([xshift=-1.5cm,yshift=-.4cm]b2.south) {};
    \node[anchor=north,text width=5.5cm,align=center] (b31) at ([yshift=-.1cm]b3.north) {Student-generated partial trajectory ($y_{<i}$)};
    \path ([xshift=.2cm,yshift=.2cm]b3.south west) node [draw,fill=gray!10,anchor=south west,minimum height=.6cm] (b32) {$3x+7=22$}
    ([xshift=1cm]b32.east) node [draw,fill=gray!10,anchor=west,minimum height=.6cm] (b33) {${3x=15}$}
    ([xshift=-.2cm,yshift=.2cm]b3.south east) node [draw,dashed,anchor=south east,minimum height=.6cm,minimum width=.6cm] (b34) {?};
    \draw[->] ([xshift=.2cm]b32.east) -- ([xshift=-.2cm]b33.west);
    \draw[->,orange,rounded corners] ([yshift=.5cm]b34.north) -- +(0,0.4cm) -- ([yshift=.9cm]b34.north-|b25.south) -- ([yshift=-.1cm]b25.south);

    \path ([xshift=-2cm,yshift=0.3cm]b1.north) node[draw=blue,dashed,rounded corners,anchor=south,minimum width=7cm, minimum height=1cm] (b4)
    {\small$R(x,y_{<i})=\sum\limits_{t=i}^nD_{\textrm{KL}}(P^t(\cdot|x,y_{<t})||P_{\theta}^s(\cdot|x,y_{<t}))$}
    (b4.north west) node [anchor=south west] (b41) {\textcolor{blue}{Computing rewards:}};
    \draw [->,blue,dashed,rounded corners] (b4.west) -- (b4-|sm) -- (sm);

    \node [text width=4cm,anchor=south west,align=left] (st) at ([xshift=2cm]b2.south east) {Student's predicted \\ Distribution $P_{\theta}^s(\cdot|x,y_{<t})$};
    \node [text width=4cm,anchor=south west,align=left] (tt) at ([yshift=.3cm]st.north west) {Teacher's predicted \\ Distribution $P^t(\cdot|x,y_{<t})$};
    \node [draw=orange,fill=orange!10,minimum width=.5cm,minimum height=.3cm,anchor=north east] at ([xshift=-.1cm,yshift=-.1cm]st.north west) {};
    \node [draw=violet,fill=violet!10,minimum width=.5cm,minimum height=.3cm,anchor=north east] at ([xshift=-.1cm,yshift=-.1cm]tt.north west) {};

  \end{tikzpicture}
\end{center}
