Add support for Craftax

.gitignore

@@ -146,4 +146,9 @@ envpool
 .sokoban_cache
 .build
 .vscode
-k8s_copy/
+k8s_copy/
+Craftax_Baselines
+
+nsight-profile
+
+uv.lock

.gitmodules

@@ -4,3 +4,6 @@
 [submodule "third_party/gym-sokoban"]
 	path = third_party/gym-sokoban
 	url = https://github.com/AlignmentResearch/gym-sokoban
+[submodule "third_party/craftax"]
+	path = third_party/craftax
+	url = https://github.com/rhaps0dy/Craftax

.pre-commit-config.yaml

@@ -2,7 +2,7 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.1.13
+  rev: v0.9.7
   hooks:
     # Run the linter.
     - id: ruff

CLAUDE.md

@@ -0,0 +1,30 @@
+# CLAUDE.md: Development Guidelines
+
+## Build & Test Commands
+- Install: `make local-install`
+- Lint: `make lint`
+- Format: `make format`
+- Typecheck: `make typecheck`
+- Run tests: `make mactest` or `pytest -m 'not envpool and not slow'`
+- Run single test: `pytest tests/test_file.py::test_function -v`
+- Training command: `python -m cleanba.cleanba_impala --from-py-fn=cleanba.config:sokoban_drc33_59`
+
+## Code Style Guidelines
+- **Formatting**: Follow ruff format (127 line length)
+- **Imports**: Use isort through ruff, with known third-party libraries like wandb
+- **Types**: Use type annotations, checked with pyright
+- **Naming**: 
+  - Variables: `snake_case`
+  - Classes: `PascalCase`
+  - Constants: `UPPER_CASE`
+- **Structure**: Modules organized by functionality (cleanba, experiments, tests)
+- **Error handling**: Use asserts for validation in tests, exceptions for runtime errors
+- **Documentation**: Include docstrings for public functions and classes
+- **JAX/Flax patterns**: Use pure functions and maintain functional style
+
+## Key Files
+- **environments.py**: Contains environment wrappers and config classes for different environments (Sokoban, Boxoban, Craftax). Includes `EpisodeEvalWrapper` for logging episode returns and adapters for different environment backends.
+
+- **cleanba_impala.py**: Main training loop implementation for the IMPALA algorithm. Contains multi-threaded rollout data collection, parameter synchronization, and training. Uses `WandbWriter` for logging, queues for communication between rollout and learner threads, and implements checkpointing.
+
+- **impala_loss.py**: Implements V-trace and PPO loss functions. Contains `Rollout` data structure, TD-error computation with V-trace, and policy gradient calculations. Handles truncated episodes specially to provide correct advantage estimates.

Dockerfile

@@ -1,36 +1,42 @@
 ARG JAX_DATE
 
-FROM ghcr.io/nvidia/jax:base-${JAX_DATE} as envpool-environment
+FROM ghcr.io/nvidia/jax:base-${JAX_DATE} AS envpool-environment
 ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update \
-    && apt-get install -y golang-1.18 git \
+    && apt-get install -y golang-1.21 git \
     # Linters
       clang-format clang-tidy \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 
-ENV PATH=/usr/lib/go-1.18/bin:/root/go/bin:$PATH
-RUN go install github.com/bazelbuild/[email protected] && ln -sf $HOME/go/bin/bazelisk $HOME/go/bin/bazel
-RUN go install github.com/bazelbuild/buildtools/[email protected]
+USER ubuntu
+ENV HOME=/home/ubuntu
+ENV PATH=/usr/lib/go-1.21/bin:${HOME}/go/bin:$PATH
+ENV UID=1000
+ENV GID=1000
+RUN --mount=type=cache,target=${HOME}/.cache,uid=${UID},gid=${GID} go install github.com/bazelbuild/[email protected] && ln -sf $HOME/go/bin/bazelisk $HOME/go/bin/bazel
+RUN --mount=type=cache,target=${HOME}/.cache,uid=${UID},gid=${GID} go install github.com/bazelbuild/buildtools/[email protected]
 # Install Go linting tools
-RUN go install github.com/google/[email protected]
+RUN --mount=type=cache,target=${HOME}/.cache,uid=${UID},gid=${GID} go install github.com/google/[email protected]
 
-ENV USE_BAZEL_VERSION=6.4.0
+ENV USE_BAZEL_VERSION=8.1.0
 RUN bazel version
 
 WORKDIR /app
+# Copy the whole repository
+COPY --chown=ubuntu:ubuntu third_party/envpool .
+
 # Install python-based linting dependencies
-COPY third_party/envpool/third_party/pip_requirements/requirements-devtools.txt \
+COPY --chown=ubuntu:ubuntu \
+    third_party/envpool/third_party/pip_requirements/requirements-devtools.txt \
     third_party/pip_requirements/requirements-devtools.txt
-RUN pip install -r third_party/pip_requirements/requirements-devtools.txt
-
-# Copy the whole repository
-COPY third_party/envpool .
+RUN --mount=type=cache,target=${HOME}/.cache,uid=1000,gid=1000 pip install -r third_party/pip_requirements/requirements-devtools.txt
+ENV PATH="${HOME}/.local/bin:${PATH}"
 
 # Deal with the fact that envpool is a submodule and has no .git directory
 RUN rm .git
 # Copy the .git repository for this submodule
-COPY .git/modules/envpool ./.git
+COPY --chown=ubuntu:ubuntu .git/modules/envpool ./.git
 # Remove config line stating that the worktree for this repo is elsewhere
 RUN sed -e 's/^.*worktree =.*$//' .git/config > .git/config.new && mv .git/config.new .git/config
 
@@ -39,14 +45,14 @@ RUN echo "$(git status --porcelain --ignored=traditional)" \
     && if ! { [ -z "$(git status --porcelain --ignored=traditional)" ] \
     ; }; then exit 1; fi
 
-FROM envpool-environment as envpool
-RUN make bazel-release
+FROM envpool-environment AS envpool
+RUN --mount=type=cache,target=${HOME}/.cache,uid=1000,gid=1000 make bazel-release && cp bazel-bin/*.whl .
 
-FROM ghcr.io/nvidia/jax:jax-${JAX_DATE} as main-pre-pip
+FROM ghcr.io/nvidia/jax:jax-${JAX_DATE} AS main-pre-pip
 
 ARG APPLICATION_NAME
-ARG USERID=1001
-ARG GROUPID=1001
+ARG UID=1001
+ARG GID=1001
 ARG USERNAME=dev
 
 ENV GIT_URL="https://github.com/AlignmentResearch/${APPLICATION_NAME}"
@@ -84,8 +90,8 @@ ENV VIRTUAL_ENV="/opt/venv"
 ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"
 
 RUN python3 -m venv "${VIRTUAL_ENV}" --system-site-packages \
-    && addgroup --gid ${GROUPID} ${USERNAME} \
-    && adduser --uid ${USERID} --gid ${GROUPID} --disabled-password --gecos '' ${USERNAME} \
+    && addgroup --gid ${GID} ${USERNAME} \
+    && adduser --uid ${UID} --gid ${GID} --disabled-password --gecos '' ${USERNAME} \
     && usermod -aG sudo ${USERNAME} \
     && echo "${USERNAME} ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers \
     && mkdir -p "/workspace" \
@@ -96,32 +102,38 @@ WORKDIR "/workspace"
 # Get a pip modern enough that can resolve farconf
 RUN pip install "pip ==24.0" && rm -rf "${HOME}/.cache"
 
-FROM main-pre-pip as main-pip-tools
+FROM main-pre-pip AS main-pip-tools
 RUN pip install "pip-tools ~=7.4.1"
 
-FROM main-pre-pip as main
+FROM main-pre-pip AS main
+RUN --mount=type=cache,target=${HOME}/.cache,uid=${UID},gid=${GID} pip install uv
 COPY --chown=${USERNAME}:${USERNAME} requirements.txt ./
 # Install all dependencies, which should be explicit in `requirements.txt`
-RUN pip install --no-deps -r requirements.txt \
-    && rm -rf "${HOME}/.cache" \
+RUN --mount=type=cache,target=${HOME}/.cache,uid=${UID},gid=${GID} \
+    uv pip install --no-deps -r requirements.txt \
     # Run Pyright so its Node.js package gets installed
     && pyright .
 
 # Install Envpool
-ENV ENVPOOL_WHEEL="dist/envpool-0.8.4-cp310-cp310-linux_x86_64.whl"
-COPY --from=envpool --chown=${USERNAME}:${USERNAME} "/app/${ENVPOOL_WHEEL}" "./${ENVPOOL_WHEEL}"
-RUN pip install "./${ENVPOOL_WHEEL}" && rm -rf "./dist"
+ENV ENVPOOL_WHEEL="envpool-0.9.0-cp312-cp312-linux_x86_64.whl"
+COPY --from=envpool --chown=${USERNAME}:${USERNAME} "/app/${ENVPOOL_WHEEL}" "${ENVPOOL_WHEEL}"
+RUN uv pip install "${ENVPOOL_WHEEL}" && rm "${ENVPOOL_WHEEL}"
+
+# Cache Craftax textures
+RUN python -c "import craftax.craftax.constants"
 
 # Copy whole repo
 COPY --chown=${USERNAME}:${USERNAME} . .
-RUN pip install --no-deps -e . -e ./third_party/gym-sokoban/
+RUN --mount=type=cache,target=${HOME}/.cache,uid=${UID},gid=${GID} \
+    uv pip install --no-deps -e . -e ./third_party/gym-sokoban/
 
 # Set git remote URL to https for all sub-repos
 RUN git remote set-url origin "$(git remote get-url origin | sed 's|[email protected]:|https://github.com/|' )" \
     && (cd third_party/envpool && git remote set-url origin "$(git remote get-url origin | sed 's|[email protected]:|https://github.com/|' )" )
 
 # Abort if repo is dirty
-RUN echo "$(git status --porcelain --ignored=traditional | grep -v '.egg-info/$')" \
+RUN rm NVIDIA_Deep_Learning_Container_License.pdf \
+    && echo "$(git status --porcelain --ignored=traditional | grep -v '.egg-info/$')" \
     && echo "$(cd third_party/envpool && git status --porcelain --ignored=traditional | grep -v '.egg-info/$')" \
     && echo "$(cd third_party/gym-sokoban && git status --porcelain --ignored=traditional | grep -v '.egg-info/$')" \
     && if ! { [ -z "$(git status --porcelain --ignored=traditional | grep -v '.egg-info/$')" ] \
@@ -130,5 +142,5 @@ RUN echo "$(git status --porcelain --ignored=traditional | grep -v '.egg-info/$'
     ; }; then exit 1; fi
 
 
-FROM main as atari
-RUN pip uninstall -y envpool && pip install envpool && rm -rf "${HOME}/.cache"
+FROM main AS atari
+RUN uv pip uninstall -y envpool && uv pip install envpool && rm -rf "${HOME}/.cache"

Makefile

@@ -8,7 +8,8 @@ export DOCKERFILE
 
 COMMIT_HASH ?= $(shell git rev-parse HEAD)
 BRANCH_NAME ?= $(shell git branch --show-current)
-JAX_DATE=2024-04-08
+JAX_DATE=2025-02-22
+PYTHON_VERSION=3.12
 
 default: release/main
 
@@ -29,21 +30,19 @@ BUILD_PREFIX ?= $(shell git rev-parse --short HEAD)
 		-f "${DOCKERFILE}" .
 	touch ".build/with-reqs/${BUILD_PREFIX}/$*"
 
-# NOTE: --extra=extra is for stable-baselines3 testing.
-requirements.txt.new: pyproject.toml ${DOCKERFILE}
-	docker run -v "${HOME}/.cache:/home/dev/.cache" -v "$(shell pwd):/workspace" "ghcr.io/nvidia/jax:base-${JAX_DATE}" \
-    bash -c "pip install pip-tools \
-		&& cd /workspace \
-		&& pip-compile --verbose -o requirements.txt.new --extra=dev --extra=launch_jobs pyproject.toml"
+requirements.txt.new: pyproject.toml
+	docker run -v "${HOME}/.cache:/home/dev/.cache" -v "$(shell pwd):/workspace" "ghcr.io/astral-sh/uv:python${PYTHON_VERSION}-alpine" \
+    sh -c "cd /workspace \
+		&& uv pip compile --verbose -o requirements.txt.new --extra=dev pyproject.toml"
 
 # To bootstrap `requirements.txt`, comment out this target
 requirements.txt: requirements.txt.new
 	sed -E "s/^(jax==.*|jaxlib==.*|nvidia-.*|torchvision==.*|torch==.*|triton==.*)$$/# DISABLED \\1/g" requirements.txt.new > requirements.txt
 
 .PHONY: local-install
 local-install: requirements.txt
-	pip install --no-deps -r requirements.txt
-	pip install --config-settings editable_mode=compat -e ".[dev-local]" -e ./third_party/gym-sokoban
+	uv pip install --no-deps -r requirements.txt
+	uv pip install -e ".[py-tools]" -e ./third_party/gym-sokoban
 	pip install https://github.com/AlignmentResearch/envpool/releases/download/v0.1.0/envpool-0.8.4-cp310-cp310-linux_x86_64.whl
 
 
@@ -93,18 +92,18 @@ cuda-devbox/%: devbox/%
 cuda-devbox: cuda-devbox/main
 
 .PHONY: envpool-devbox
-envpool-devbox: devbox/envpool-ci
+envpool-devbox: devbox/envpool
 
 
 .PHONY: docker docker/%
 docker/%:
-	docker run -v "$(shell pwd):/workspace" -it "${APPLICATION_URL}:${RELEASE_PREFIX}-$*" /bin/bash
+	docker run -v "${HOME}/.cache:/home/ubuntu/.cache" -v "$(shell pwd):/workspace" -it "${APPLICATION_URL}:${RELEASE_PREFIX}-$*" /bin/bash
 docker: docker/main
 
 .PHONY: envpool-docker envpool-docker/%
 envpool-docker/%:
-	docker run -v "$(shell pwd)/third_party/envpool:/app" -it "${APPLICATION_URL}:${RELEASE_PREFIX}-$*" /bin/bash
-envpool-docker: envpool-docker/envpool-ci
+	docker run -v "${HOME}/.cache:/home/ubuntu/.cache" -v "$(shell pwd)/third_party/envpool:/app" -it "${APPLICATION_URL}:${RELEASE_PREFIX}-$*" /bin/bash
+envpool-docker: envpool-docker/envpool
 
 # Section 3: project commands