Update figure of opd in section5

section5/Figures/opd_illustration.tex

+\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}

section5/section5.tex

@@ -137,8 +137,9 @@ Compared with conventional offline distillation, OPD adopts a simple but importa
 
 \begin{figure*}[!t]
     \centering
-    \resizebox{\linewidth}{!}{
-    \includegraphics[width=0.5\textwidth]{section5/Figures/opd.png}}
+    \input{section5/Figures/opd_illustration.tex}
+    % \resizebox{\linewidth}{!}{
+    % \includegraphics[width=0.5\textwidth]{section5/Figures/opd.png}}
     \caption{
     Illustration of OPD \citep{agarwal-etal:onpolicy}. The student model generates trajectories from its current policy, while the teacher model provides next-token distributional supervision at the states visited by the student. The divergence between the teacher and student distributions is used to construct dense token-level rewards for optimizing the student model.}
     \label{fig:opd-implementation}