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

% Extracted memory
\node[box, text width=2.8cm] (extracted) {
\textbf{Extracted Memory}\\[0.1cm]
$x^{\mathrm{mem}}$
};

% Memory manager
\node[box, text width=3.2cm, right=1.5cm of extracted] (manager) {
\textbf{Memory Manager}\\[0.1cm]
$\pi_\theta$\\
Select memory operation
};

% Memory bank
\node[box, text width=3.2cm, 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 manager] (agent) {
\textbf{Agent}\\[0.1cm]
Downstream task solving
};

% Reward
\node[box, text width=2.8cm, below=1.4cm of agent] (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=\small]above:Learnable Memory Management}
] (managerbox) {};

% Arrows
\draw[arrow] (extracted) -- node[above,font=\small] {input} (manager);

\draw[arrow] (manager) -- node[right,font=\small] {operation} (bank);

\draw[arrow] (bank.east) -- ++(0.8,0)
    |- node[pos=0.25,above,font=\small] {retrieve}
    (agent.west);

\draw[arrow] (agent.west) -- ++(-0.8,0)
    |- node[pos=0.25,below,font=\small] {updated memory}
    (bank.east);

\draw[reward] (agent.south) -- node[right,font=\small] {feedback} (rewardbox.north);

\draw[reward] (rewardbox.west) -- ++(-1.5,0)
    |- node[pos=0.25,left,font=\small] {RL optimization}
    (manager.south);

\end{tikzpicture}