Learning Memory Mechanisms for Decision Making through Demonstration

In Partially Observable Markov Decision Processes, integrating an agent's history into memory poses a significant challenge for decision-making. Traditional imitation learning, relying on observation-action pairs for expert demonstrations, fails to capture the expert's memory mechanisms used in decision-making. To capture memory processes as demonstrations, we introduce the concept of **memory dependency pairs** $(p, q)$ indicating that events at time $p$ are recalled for decision-making at time $q$. We introduce **AttentionTuner** to leverage memory dependency pairs in Transformers and find significant improvements across several tasks compared to standard Transformers when evaluated on Memory Gym and the Long-term Memory Benchmark.