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

\begin{center}
    \begin{tikzpicture}
      \tikzstyle{model} = [minimum width=2cm, minimum height=1.2cm,align=center,inner sep=2pt,draw,thick,drop shadow={fill=gray,shadow xshift=.6ex,shadow yshift=-.6ex},fill=white]
      \tikzstyle{label} = [minimum width=2ex, minimum height=1.5ex,draw,scale=.75,fill=white]
      \tikzstyle{circled} = [shape=circle,draw,inner sep=2pt]
      \begin{scope}
        \node [anchor=west] (x) at (-0.2cm,0) {$\mathbf{x}$};
        \node [anchor=west,model] (policy model) at ([xshift=.8cm]x) {Policy\\Model};
        \node [anchor=north west,label] at ([xshift=-.1cm,yshift=.15cm]policy model.north west) {To Learn};
        \node [anchor=west,align=center] (y) at ([xshift=1.2cm]policy model.east) {$\mathbf{y}_1$\\[1ex]$\mathbf{y}_2$\\[1ex]$\cdots$\\[1ex]$\mathbf{y}_G$};
        \node [anchor=west,minimum width=2cm, minimum height=1.2cm] (reward model) at ([xshift=1.4cm]y.east) {};
        \node [anchor=south,model] (reference model) at ([yshift=.4cm]reward model.north) {Reference\\Model};
        \node [anchor=north west,label] at ([xshift=-.1cm,yshift=.15cm]reference model.north west) {Fixed};
        \node [anchor=north,model] (value model) at ([yshift=-.4cm]reward model.south) {Reward\\Model};
        \node [anchor=north west,label] at ([xshift=-.1cm,yshift=.15cm]value model.north west) {Fixed};

        \node [anchor=west] (reward) at ([xshift=0.1cm]reward model.east) {};
        \node [anchor=south west,align=center] (value) at ([xshift=0.1cm,yshift=-.5cm]value model.south east) {$R_\phi(\mathbf{x}, \mathbf{y}_1)$\\[1ex]$R_\phi(\mathbf{x}, \mathbf{y}_2)$\\[1ex]$\cdots$\\[1ex]$R_\phi(\mathbf{x}, \mathbf{y}_G)$};

        \node [anchor=west,align=center,rounded corners=5pt,inner sep=12pt,fill=lightgray] (advantage estimation) at ([xshift=1.5cm]reward) {Advantage\\Estimation};
        \node [anchor=west] (advantage estimation out) at (advantage estimation.east) {$\{A^{\mathrm{grpo}}_{1,t},\cdots, A^{\mathrm{grpo}}_{G,t}\}$};
        

        \node (left bottom) at (x.west|-value model.south) {};
        \node (right bottom) at (advantage estimation out.east|-value model.south) {};
        \node [anchor=north,minimum width=13.6cm, minimum height=2cm,rounded corners=8pt,draw=lightgray,thick] (note) at ($(left bottom)!.5!(right bottom)+(0, -1cm)$) {};
        \matrix [anchor=center, every even column/.style={column sep=1cm}] at (note.center) {
        \node {\circled{1}}; & \node [anchor=west, inner sep=0pt] {sample multiple outputs $\{\mathbf{y}_{1}, \mathbf{y}_{2}, \cdots, \mathbf{y}_{G}\}$} ; & \node {\circled{2}}; & \node [anchor=west, inner sep=0pt] {compute the $\mathrm{Penalty}$} ; \\
        \node {\circled{3}}; & \node [anchor=west, inner sep=0pt] {compute the reward for all sampled outputs} ; & \node {\circled{4}}; & \node [anchor=west, inner sep=0pt] {optimize the policy model} ; \\
        };
        
        \draw [->,thick] (x.east) -- (policy model.west);
        \draw [->,thick] (policy model.east) node [xshift=.5cm,yshift=.0cm,above,align=center] {\circled{1}} -- (y.west);
        \draw [->,thick] (y.east)  .. controls (reward model.west) and ($(y.east |- reference model.west)$) .. (reference model.west);
        \draw [->,thick] (y.east) .. controls (reward model.west) and ($(y.east |- value model.west)$) .. (value model.west);

        \draw [->,thick] ([xshift=2cm]value model.east) .. controls (value model.east -| advantage estimation.south) .. node [above,xshift=-.2cm] {\circled{3}}(advantage estimation.south);
        \draw [->,thick, dashed] (reference model.north) .. controls ([yshift=.5cm]reference model.north) .. ([xshift=-1cm, yshift=.5cm]reference model.north) node (tmp) {} -- node [midway, above] {$\mathrm{Penalty}$} ([xshift=1cm]policy model.north|-tmp) .. controls (policy model.north|-tmp) .. ([yshift=-1cm]policy model.north|-tmp) -- node [right, solid] {\circled{2}} ([yshift=.2cm]policy model.north);

        % update model
        \node (policy model bottom) at ($(right bottom -| policy model)+(0,-.6cm)$) {};
        \node [anchor=north] (policy model bottom1) at (policy model bottom |- advantage estimation out.north) {};
        \node (value model bottom) at ($(right bottom -| value model)+(0,-.6cm)$) {};
        \node (advantage estimation out bottom) at ($(advantage estimation out |- right bottom)+(0,-.6cm)$) {};
        \draw [->,thick,dashed] (advantage estimation out) -- ($(advantage estimation out)!.5!(advantage estimation out bottom)$) .. controls (advantage estimation out bottom) .. ($(advantage estimation out bottom.center)!.5!(policy model bottom.center)+(5em,0)$) -- ($(advantage estimation out bottom.center)!.5!(policy model bottom.center)-(5em,0)$)  .. controls (policy model bottom.center) .. ($(policy model bottom1.center)!.5!(policy model bottom)$) -- node [above,xshift=0.3cm,yshift=-.5cm,solid] {\circled{4}}(policy model.south);

      \end{scope}
    \end{tikzpicture}    
\end{center}
