This repository contains the official implementation of Delay-Aware Dreamer (DA-Dreamer), a model-based reinforcement learning framework for random observation delays in POMDPs, introduced in our paper. The code is built on top of DreamerV3 (JAX) and extends it with a latent-space filtering mechanism** that enables robust decision making under stochastic and out-of-sequence observation delays. Please read the paper for more details.
Create a Conda environment using the provided configuration:
conda env create -n dadreamer -f environment.yml
conda activate dadreamerExample commands:
- The main method (referred as DA-Dreamer) with uniform delays on a Gym environment:
python dreamerv3/main.py --logdir ./logdir/$(date "+%Y%m%d-%H%M%S")\
--configs gym_proprio latent size7m\
--task gym_HalfCheetah-v4\
--run.steps 1e6\
--delay.obs_delay_step False\
--replay.save_replay False\
--agent.repval_loss False\
--delay.maximum_delay 20 --delay.obs_window_size 20\
--delay.delay_distribution Uniform\
--delay.delay_Uniform_dist_kwargs.minimum 0\
--delay.delay_Uniform_dist_kwargs.maximum 20- Treating delays as a generic input (referred as Stack-Dreamer) with TruncatedGaussian on a Gym environment:
python dreamerv3/main.py --logdir ./logdir/$(date "+%Y%m%d-%H%M%S")\
--configs gym_proprio extended_obs size7m\
--task gym_HalfCheetah-v4\
--run.steps 1e6\
--delay.privilaged_decoder False\
--agent.enc.simple.enc_key_prefix ''\
--agent.dec.simple.dec_key_prefix ''\
--delay.include_actions True\
--replay.save_replay False\
--agent.repval_loss False\
--delay.maximum_delay 40 --delay.obs_window_size 40\
--delay.delay_distribution TruncatedGaussian\
--delay.delay_TruncatedGaussian_dist_kwargs.minimum 0\
--delay.delay_TruncatedGaussian_dist_kwargs.maximum 40\
--delay.delay_TruncatedGaussian_dist_kwargs.mu 10\
--delay.delay_TruncatedGaussian_dist_kwargs.scale 1- Wait agent on a Meta-Word environment with added action noise. Replace
waitwithlatest_obsto switch to Memoryless agent:
python dreamerv3/main.py --logdir ./logdir/$(date "+%Y%m%d-%H%M%S")\
--configs metaworld_vision wait size1m log_tb\
--task metaworld_drawer-close-v2\
--run.steps 5e5\
--env.metaworld.length 50\
--delay.obs_delay_step True\
--replay.save_replay False\
--env.action_noise.noise_fraction 0.2\
--agent.repval_loss False\
--delay.maximum_delay 5 --delay.obs_window_size 5\
--delay.delay_distribution Uniform\
--delay.delay_Uniform_dist_kwargs.minimum 0\
--delay.delay_Uniform_dist_kwargs.maximum 5To set the number of particles, use --agent.dyn.rssm.particles number_of_particles.
If you find this code useful, please reference DreamerV3 and our paper:
@article{karamzade2025model,
title={Model-Based Reinforcement Learning under Random Observation Delays},
author={Karamzade, Armin and Kim, Kyungmin and Lanier, JB and Corsi, Davide and Fox, Roy},
journal={arXiv preprint arXiv:2509.20869},
year={2025}
}We would like to thank the authors of DreamerV3 for publicly sharing their work.