% !Mode:: "TeX:UTF-8"
% !TEX encoding = UTF-8 Unicode

\begin{center}
  \begin{tikzpicture}
    
    \def\ssep{0.9cm}

    \tikzstyle{lnode} = [minimum width=3cm,minimum height=1.0cm,draw,thick,fill=white,inner sep=2pt,drop shadow={fill=gray,shadow xshift=.6ex,shadow yshift=-.6ex}]; % rounded corners=1pt,

    %%% RLHF
    \begin{scope}

      \node [anchor=west,fill=gray!30,inner sep=5pt] (preferencedata) at (0,0) {\Large{$\mathbf{y}_a \succ \mathbf{y}_b$}};
      \node [anchor=south west,align=left] (preferencelabel) at ([yshift=0.1cm]preferencedata.north west) {\scriptsize{Preference}\\[-1mm] \scriptsize{Data}};

      \node [lnode,anchor=west] (rewardmodel) at ([xshift=2.5cm]preferencedata.east) {Reward Model};

      \node [anchor=north west,minimum height=5.8cm,minimum width=3.8cm,rounded corners=5pt,fill=lightgray] (bg) at ([xshift=2.1cm,yshift=2.9cm]rewardmodel.east) {};
      
      \node [lnode,anchor=south west] (valuefunction) at ([xshift=2.5cm,yshift=1.4cm]rewardmodel.east) {Value Model};

      \node [lnode,anchor=west] (policy) at ([xshift=2.5cm,yshift=-0.0cm]rewardmodel.east) {Policy Model};

      \node [lnode,anchor=north west] (reference) at ([xshift=2.5cm,yshift=-1.4cm]rewardmodel.east) {Reference Model};

      \node [align=left] at ([yshift=-1.35cm]reference) {\scriptsize{Training with PPO}};

      \draw [->,thick] ([xshift=2pt]preferencedata.east)--([xshift=-2pt]rewardmodel.west) node [pos=0.5,above,align=left,xshift=-0.05cm] {\scriptsize{training with MLE}};

      \draw [->,thick] ([xshift=2pt]rewardmodel.east)--([xshift=-2pt]bg.west);

      % \draw [->,thick] (rewardmodel.east) .. controls ($(rewardmodel.east)!.5!(policy.west)$) and (rewardmodel.east|-valuefunction.west) .. (valuefunction.west);
      % \draw [->,thick] (rewardmodel.east) -- (policy.west);
      % \draw [->,thick] (rewardmodel.east) .. controls ($(rewardmodel.east)!.5!(policy.west)$) and (rewardmodel.east|-reference.west) .. (reference.west);

      \draw [->,dotted,very thick] ([xshift=-0.6cm,yshift=-2pt]valuefunction.south) .. controls +(-150:0.5cm) and +(150:0.5cm) .. ([xshift=-0.6cm,yshift=2pt]policy.north);
      \draw [<-,dotted,very thick] ([xshift=0.6cm,yshift=-2pt]valuefunction.south) .. controls +(-30:0.5cm) and +(30:0.5cm) .. ([xshift=0.6cm,yshift=2pt]policy.north);
      \draw [<-,dotted,very thick] ([xshift=-0.6cm,yshift=+2pt]reference.north) .. controls +(150:0.5cm) and +(-150:0.5cm) .. ([xshift=-0.6cm,yshift=-2pt]policy.south);
      \draw [->,dotted,very thick] ([xshift=0.6cm,yshift=+2pt]reference.north) .. controls +(30:0.5cm) and +(-30:0.5cm) .. ([xshift=0.6cm,yshift=-2pt]policy.south);

      \node (mid) at ($(preferencedata.west)!.5!(bg.east)$) {};
      \node [anchor=north] (caption) at ([yshift=-.5cm]mid.center|-bg.south) {\small{(a) Training LLMs with PPO}};

    \end{scope}

    %%% DPO
    \begin{scope}[yshift=-5.5cm]
    \node [anchor=west,fill=gray!30,inner sep=5pt] (preferencedata) at (0,0) {\Large{$\mathbf{y}_a \succ \mathbf{y}_b$}};
      \node [anchor=south west,align=left] (preferencelabel) at ([yshift=0.1cm]preferencedata.north west) {\scriptsize{Preference}\\ [-1mm] \scriptsize{Data}};

      \node [anchor=north west,minimum height=2cm,minimum width=3.8cm,rounded corners=5pt,fill=lightgray] (bg2) at ([xshift=7.6cm,yshift=1cm]preferencedata.east) {};
      \node [lnode,anchor=west] (policy) at ([xshift=8cm]preferencedata.east) {Policy Model};

      \draw [->,thick] ([xshift=2pt]preferencedata.east)--([xshift=-2pt]bg2.west) node [pos=0.5,above,align=left] {\scriptsize{training with MLE}};

      \node (mid2) at ($(preferencedata.west)!.5!(bg2.east)$) {};
      \node [anchor=north] (caption b) at ([xshift=-0.0cm,yshift=-.18cm]mid2.center|-bg2.south) {\small{(b) Training LLMs with DPO}};

    \end{scope}

  \end{tikzpicture}
\end{center}
