Move SharedAdam to utils

Author: simeetnayan81

.github/scripts/pre-build-script-win.sh

@@ -2,4 +2,4 @@
 
 pip install --upgrade setuptools
 
-export TORCHRL_BUILD_VERSION=0.9.0
+export TORCHRL_BUILD_VERSION=0.10.0

.github/scripts/td_script.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-export TORCHRL_BUILD_VERSION=0.9.0
+export TORCHRL_BUILD_VERSION=0.10.0
 pip install --upgrade setuptools
 
 # Check if ARCH is set to aarch64

.github/scripts/version_script.bat

@@ -1,5 +1,5 @@
 @echo off
-set TORCHRL_BUILD_VERSION=0.9.0
+set TORCHRL_BUILD_VERSION=0.10.0
 echo TORCHRL_BUILD_VERSION is set to %TORCHRL_BUILD_VERSION%
 
 @echo on

.github/unittest/linux/scripts/environment.yml

@@ -35,3 +35,4 @@ dependencies:
     - transformers
     - ninja
     - timm
+    - safetensors

.github/unittest/linux_libs/scripts_brax/environment.yml

@@ -21,3 +21,4 @@ dependencies:
     - hydra-core
     - jax[cuda12]
     - brax
+    - psutil

.github/unittest/linux_libs/scripts_brax/run_test.sh

@@ -8,6 +8,13 @@ conda activate ./env
 
 export PYTORCH_TEST_WITH_SLOW='1'
 export LAZY_LEGACY_OP=False
+
+# Configure JAX for proper GPU initialization
+export XLA_PYTHON_CLIENT_PREALLOCATE=false
+export XLA_PYTHON_CLIENT_ALLOCATOR=platform
+export TF_FORCE_GPU_ALLOW_GROWTH=true
+export CUDA_VISIBLE_DEVICES=0
+
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
@@ -28,7 +35,33 @@ export MAGNUM_LOG=verbose MAGNUM_GPU_VALIDATION=ON
 # this workflow only tests the libs
 python -c "import brax"
 python -c "import brax.envs"
-python -c "import jax"
+
+# Initialize JAX with proper GPU configuration
+python -c "
+import jax
+import jax.numpy as jnp
+import os
+
+# Configure JAX for GPU
+os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
+os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = 'platform'
+
+# Test JAX GPU availability
+try:
+    devices = jax.devices()
+    print(f'JAX devices: {devices}')
+    if len(devices) > 1:
+        print('JAX GPU is available')
+    else:
+        print('JAX CPU only')
+except Exception as e:
+    print(f'JAX initialization error: {e}')
+    # Fallback to CPU
+    os.environ['JAX_PLATFORM_NAME'] = 'cpu'
+    jax.config.update('jax_platform_name', 'cpu')
+    print('Falling back to JAX CPU')
+"
+
 python3 -c 'import torch;t = torch.ones([2,2], device="cuda:0");print(t);print("tensor device:" + str(t.device))'
 
 python .github/unittest/helpers/coverage_run_parallel.py -m pytest test/test_libs.py --instafail -v --durations 200 --capture no -k TestBrax --error-for-skips

.github/unittest/linux_libs/scripts_minari/environment.yml

@@ -20,4 +20,11 @@ dependencies:
     - scipy
     - hydra-core
     - minari[gcs,hdf5,hf]
-    - gymnasium<1.0.0
+    - gymnasium>=1.2.0
+    - ale-py
+    - gymnasium-robotics
+    - minari[create]
+    - jax
+    - mujoco               
+    - mujoco-py<2.2,>=2.1
+    - minigrid

.github/unittest/linux_libs/scripts_minari/install.sh

@@ -7,8 +7,7 @@ unset PYTORCH_VERSION
 
 set -e
 
-eval "$(./conda/bin/conda shell.bash hook)"
-conda activate ./env
+# Note: This script is sourced by run_all.sh, so the environment is already active
 
 if [ "${CU_VERSION:-}" == cpu ] ; then
     version="cpu"
@@ -22,22 +21,21 @@ else
     version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
 fi
 
-
 # submodules
 git submodule sync && git submodule update --init --recursive
 
 printf "Installing PyTorch with cu128"
 if [[ "$TORCH_VERSION" == "nightly" ]]; then
   if [ "${CU_VERSION:-}" == cpu ] ; then
-      pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U
+      uv pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U
   else
-      pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128 -U
+      uv pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128 -U
   fi
 elif [[ "$TORCH_VERSION" == "stable" ]]; then
     if [ "${CU_VERSION:-}" == cpu ] ; then
-      pip3 install torch --index-url https://download.pytorch.org/whl/cpu
+      uv pip install torch --index-url https://download.pytorch.org/whl/cpu
   else
-      pip3 install torch --index-url https://download.pytorch.org/whl/cu128
+      uv pip install torch --index-url https://download.pytorch.org/whl/cu128
   fi
 else
   printf "Failed to install pytorch"
@@ -46,9 +44,9 @@ fi
 
 # install tensordict
 if [[ "$RELEASE" == 0 ]]; then
-  pip3 install git+https://github.com/pytorch/tensordict.git
+  uv pip install git+https://github.com/pytorch/tensordict.git
 else
-  pip3 install tensordict
+  uv pip install tensordict
 fi
 
 # smoke test

.github/unittest/linux_libs/scripts_minari/post_process.sh

@@ -2,5 +2,4 @@
 
 set -e
 
-eval "$(./conda/bin/conda shell.bash hook)"
-conda activate ./env
+# Note: This script is sourced by run_all.sh, so the environment is already active

.github/unittest/linux_libs/scripts_minari/requirements.txt

@@ -0,0 +1,20 @@
+hypothesis
+future
+cloudpickle
+pytest
+pytest-cov
+pytest-mock
+pytest-instafail
+pytest-rerunfailures
+pytest-error-for-skips
+pytest-asyncio
+expecttest
+pybind11[global]
+pyyaml
+scipy
+hydra-core
+minari[gcs,hdf5,hf,create]
+gymnasium>=1.2.0
+ale-py
+gymnasium-robotics
+mujoco