

\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)
    {\makebox[\linewidth][c]{\textbf{Extracted Memory}}\\[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]
      updating memory
    };

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

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

    % 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,align=center,font=\scriptsize] {sampling operations} (bank);
    \draw[arrow] ([yshift=0.2cm]bank.east) -- node[midway,above,font=\scriptsize,text width=1.7cm,align=center] {retrieving} ([yshift=0.2cm]agent.west);
    \draw[arrow] ([yshift=-0.2cm]agent.west) -- node[midway,below,font=\scriptsize,text width=1.7cm,align=center] {useful memory} ([yshift=-0.2cm]bank.east);
    \draw[arrow] (agent.north) -- node[right,font=\scriptsize] { trajectory} (rewardbox.south);
    \draw[arrow] (rewardbox.west) -- node[midway,above,yshift=0.08cm,inner sep=1pt,font=\scriptsize,text width=1.7cm,align=center] {RL training} (manager.east);

  \end{tikzpicture}  
\end{center}
