Add support for Craftax

[rhaps0dy]

No description provided.

[rhaps0dy]

PR Review: Add support for Craftax

This PR introduces support for Craftax, a JAX-based Minecraft-like environment for reinforcement learning. The implementation is quite extensive, with over 50 commits and changes to 29 files.

Key Changes:

1. Environment Integration:

   • Adds CraftaxEnvConfig class to support both symbolic and pixel-based environments
   • Introduces a JAX-native vectorized environment interface
   • Implements observation handling for Craftax's unique observation space

2. Infrastructure Updates:

   • Updates Dockerfile with newer dependencies (Bazel 8.1.0, Python 3.12, EnvPool 0.9.0)
   • Adds "uv" for faster Python package installation
   • Upgrades Go from 1.18 to 1.24.0

3. Training Pipeline Adjustments:

   • Refactors IMPALA and PPO implementations to work with JAX arrays
   • Adds checkpoint saving and loading functionality
   • Improves timeout handling in the training queue

4. Experiment Configuration:

   • Provides a sample experiment configuration for Craftax (DRC PPO IMPALA)
   • Configures training for 1 billion timesteps with appropriate resource allocations

Build/Test Status:

• CI is currently failing on both format and Python tests
• Format issues likely need addressing before merge
• Python test failures may indicate integration issues with existing codebase

Areas for Improvement:

1. Documentation:

   • Could benefit from more comprehensive docstrings for Craftax-specific code
   • A dedicated README or guide for using Craftax would be helpful

2. Code Quality:

   • Some hardcoded values in observation handling could be made more flexible
   • Error handling could be improved in certain areas

3. Testing:

   • Additional tests specific to Craftax functionality would strengthen the PR
   • Current test failures need to be resolved

Overall Assessment:

This is a substantial PR that adds significant new functionality to the codebase. The Craftax implementation appears well-integrated with the existing architecture, but with several architectural changes that modernize the codebase toward JAX-native operations.

Before merging, the failing CI checks should be addressed. The format failures are likely straightforward to fix, while the Python test failures may require more investigation to ensure compatibility with existing functionality.

The overall approach of using JAX for environment interaction aligns well with modern reinforcement learning practices and should provide performance benefits. The PR includes thoughtful design choices like separating symbolic and pixel-based representations and providing flexible observation transformations.