Fix bug 1

Dockerfile

@@ -24,7 +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
+ENV OMP_NUM_THREADS=4
 
 # Expose port
 EXPOSE 7860

requirements.txt

@@ -28,6 +28,15 @@ h5py>=3.7.0
 ale_py>=0.8.0
 gymnasium>=0.28.0
 
+# Experiment tracking (required by utils.py)
+wandb>=0.13.0
+
+# Metrics (required by rew_end_model.py)
+torcheval>=0.0.6
+
+# Progress bars (may be used by various components)
+tqdm>=4.64.0
+
 # Optional: for better performance
 # torch-audio  # if needed for audio processing
 

src/coroutines/env_loop.py

@@ -6,7 +6,7 @@ import torch.nn as nn
 from torch.distributions.categorical import Categorical
 
 from . import coroutine
-from envs import TorchEnv, WorldModelEnv
+from ..envs import TorchEnv, WorldModelEnv
 
 
 @coroutine

src/data/dataset.py

@@ -13,7 +13,7 @@ from torch.utils.data import Dataset as TorchDataset
 from .episode import Episode
 from .segment import Segment, SegmentId
 from .utils import make_segment
-from utils import StateDictMixin
+from ..utils import StateDictMixin
 
 
 class Dataset(StateDictMixin, TorchDataset):

src/envs/world_model_env.py

@@ -9,9 +9,9 @@ from torch import Tensor
 from torch.distributions.categorical import Categorical
 import torch.nn.functional as F
 
-from coroutines import coroutine
-from models.diffusion import Denoiser, DiffusionSampler, DiffusionSamplerConfig
-from models.rew_end_model import RewEndModel
+from ..coroutines import coroutine
+from ..models.diffusion import Denoiser, DiffusionSampler, DiffusionSamplerConfig
+from ..models.rew_end_model import RewEndModel
 
 ResetOutput = Tuple[torch.FloatTensor, Dict[str, Any]]
 StepOutput = Tuple[Tensor, Tensor, Tensor, Tensor, Dict[str, Any]]

src/models/actor_critic.py

@@ -10,9 +10,9 @@ from torch.distributions.categorical import Categorical
 import torch.nn.functional as F
 
 from .blocks import Conv3x3, SmallResBlock
-from coroutines.env_loop import make_env_loop
-from envs import TorchEnv, WorldModelEnv
-from utils import init_lstm, LossAndLogs
+from ..coroutines.env_loop import make_env_loop
+from ..envs import TorchEnv, WorldModelEnv
+from ..utils import init_lstm, LossAndLogs
 
 
 ActorCriticOutput = namedtuple("ActorCriticOutput", "logits_act val hx_cx")

src/models/diffusion/denoiser.py

@@ -10,11 +10,11 @@ from PIL import Image
 import numpy as np
 import cv2
 
-from data import Batch
+from ...data import Batch
 from .inner_model import InnerModel, InnerModelConfig
-from utils import LossAndLogs
+from ...utils import LossAndLogs
 
-from models.contour_detection_model import ContourDetectionModel
+from ..contour_detection_model import ContourDetectionModel
 
 def add_dims(input: Tensor, n: int) -> Tensor:
     return input.reshape(input.shape + (1,) * (n - input.ndim))

src/models/rew_end_model.py

@@ -8,8 +8,8 @@ import torch.nn.functional as F
 from torcheval.metrics.functional import multiclass_confusion_matrix
 
 from .blocks import Conv3x3, Downsample, ResBlocks
-from data import Batch
-from utils import init_lstm, LossAndLogs
+from ..data import Batch
+from ..utils import init_lstm, LossAndLogs
 
 
 @dataclass