add assertion for compile=true compatibility

docker/environment.yaml

@@ -14,6 +14,7 @@ dependencies:
   - torchvision
   - pip:
     - absl-py==2.1.0
+    - "cython<3"
     - dm-control==1.0.8
     - glfw==2.7.0
     - ffmpeg==1.4
@@ -23,6 +24,8 @@ dependencies:
     - hydra-core==1.3.2
     - hydra-submitit-launcher==1.2.0
     - submitit==1.5.1
+    - setuptools==65.5.0
+    - patchelf==0.17.2.1
     - omegaconf==2.3.0
     - moviepy==1.0.3
     - mujoco==2.3.1
@@ -34,13 +37,11 @@ dependencies:
     - tqdm==4.66.4
     - pandas==2.0.3
     - wandb==0.17.4
-    - matplotlib==3.7.5
-    - seaborn==0.13.2
-    - gpustat==1.1.1
+    - wheel==0.38.0
     ####################
     # Gym:
     # (unmaintained but required for maniskill2/meta-world)
-    - gym==0.21.0
+    # - gym==0.21.0
     ####################
     # ManiSkill2:
     # (requires gym==0.21.0 which occasionally breaks)

tdmpc2/common/parser.py

@@ -77,4 +77,9 @@ def parse_cfg(cfg: OmegaConf) -> OmegaConf:
 		cfg.task_dim = 0
 	cfg.tasks = TASK_SET.get(cfg.task, [cfg.task])
 
+	# Check torch.compile compatibility
+	if cfg.get('compile', False):
+		assert cfg.obs == 'state', 'torch.compile only supports state observations at the moment.'
+		assert not cfg.multitask, 'torch.compile does not support multitask training at the moment.'
+
 	return cfg_to_dataclass(cfg)