Fix bug

Dockerfile

@@ -24,6 +24,7 @@ RUN mkdir -p csgo/spawn config checkpoints cache
 # Set environment variables
 ENV PYTHONPATH=/app/src:/app
 ENV CUDA_VISIBLE_DEVICES=""
+ENV OMP_NUM_THREADS=1
 
 # Expose port
 EXPOSE 7860

requirements.txt

@@ -24,6 +24,10 @@ huggingface-hub>=0.10.0
 # Data handling
 h5py>=3.7.0
 
+# Atari environment dependencies
+ale_py>=0.8.0
+gymnasium>=0.28.0
+
 # Optional: for better performance
 # torch-audio  # if needed for audio processing
 

src/agent.py

@@ -5,11 +5,11 @@ from typing import Optional, Union
 import torch
 import torch.nn as nn
 
-from envs import TorchEnv, WorldModelEnv
-from models.actor_critic import ActorCritic, ActorCriticConfig, ActorCriticLossConfig
-from models.diffusion import Denoiser, DenoiserConfig, SigmaDistributionConfig
-from models.rew_end_model import RewEndModel, RewEndModelConfig
-from utils import extract_state_dict
+from .envs import TorchEnv, WorldModelEnv
+from .models.actor_critic import ActorCritic, ActorCriticConfig, ActorCriticLossConfig
+from .models.diffusion import Denoiser, DenoiserConfig, SigmaDistributionConfig
+from .models.rew_end_model import RewEndModel, RewEndModelConfig
+from .utils import extract_state_dict
 
 
 @dataclass