

\begin{center}
  \begin{tikzpicture}[
    box/.style={
      draw,
      fill=white,
      drop shadow={shadow xshift=0.08cm, shadow yshift=-0.08cm},
      font=\scriptsize,
      align=center,
      minimum height=1.2cm
    },
    arrow/.style={
      ->,
      thick
    }
    ]

    % Extracted memory
    \node[box, text width=3cm,align=left] (extracted)
    {\begin{center} \textbf{Extracted Memory} \end{center} \\[0.1cm]
      $\bullet$ prefers window seat\\ $\bullet$ vegetarian meals\\ $\bullet$ budget hotel};

    % Memory manager
    \node[box, text width=3cm, right=1.5cm of extracted] (manager) {
      \textbf{Memory Manager}
    };

    % Memory bank
    \node[box, text width=3cm, below=0.8cm of manager] (bank) {
      \textbf{Memory Bank}\\[0.1cm]
      $\mathcal{M}_{old}
      \rightarrow
      \mathcal{M}_{new}$
    };

    % Agent
    \node[box, text width=2.8cm, right=2.0cm of bank.east] (agent) {
      \textbf{Agent}\\[0.1cm]
      Downstream task solving
    };

    % Reward
    \node[box, text width=2.8cm, right=2.0cm of manager.east] (rewardbox) {
      \textbf{Task Feedback}\\[0.1cm]
      Reward $R$
    };

    % Dashed container for memory manager
    \node[
    draw,
    dashed,
    rounded corners,
    inner sep=0.35cm,
    fit=(manager)(bank),
    label={[font=\scriptsize]above:Learnable Memory Management}
    ] (managerbox) {};

    % Arrows
    \draw[arrow] (extracted) -- node[xshift=-.2cm,above,font=\scriptsize] {input} (manager);
    \draw[arrow] (manager) -- node[xshift=0.1cm,center,font=\scriptsize] {sampling operations} (bank);
    \draw[arrow] ([yshift=0.2cm]bank.east) -- node[above,font=\scriptsize] {retrieve} ([yshift=0.2cm]agent.west);
    \draw[arrow] ([yshift=-0.2cm]agent.west) -- node[below,font=\scriptsize,align=right] {updated\\ memory} ([yshift=-0.2cm]bank.east);
    \draw[arrow] (agent.north) -- node[right,font=\scriptsize] {feedback} (rewardbox.south);
    \draw[arrow] (rewardbox.west) -- node[xshift=0.2cm,font=\scriptsize,align=center] {RL \\optimization} (manager.east);

  \end{tikzpicture}  
\end{center}
