Merge branch 'main' into feat/total_timestep_util

.github/workflows/release.yml

@@ -7,6 +7,8 @@ on:
 jobs:
   deploy:
     runs-on: ubuntu-latest
+    container:
+      image: ghcr.io/catthehacker/ubuntu:runner-latest
     timeout-minutes: 30
     steps:
       - uses: actions/checkout@v2

.github/workflows/tests_linters.yml

@@ -4,26 +4,22 @@ on: [ push, pull_request ]
 
 jobs:
   tests-and-linters:
-    name: "Python ${{ matrix.python-version }} on ${{ matrix.os }}"
-    runs-on: "${{ matrix.os }}"
-
-    strategy:
-      matrix:
-        python-version: ["3.9"]
-        os: [self-hosted]
+    name: "Python 3.9 on GitHub Hosted runner"
+    runs-on: ubuntu-latest
+    container:
+        image: python:3.9
 
     steps:
       - name: Install dependencies for viewer test
-        run: sudo apt-get update && sudo apt-get install -y xvfb
+        run: apt-get update && apt-get install -y xvfb
       - name: Checkout flashbax
         uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
-        with:
-            python-version: "${{ matrix.python-version }}"
       - name: Install python dependencies 🔧
         run: pip install .[dev]
       - name: List python packages 📦
         run: pip list
+      - name: Update git permissions
+        run: git config --global --add safe.directory /__w/flashbax/flashbax
       - name: Run linters 🖌️
         run: pre-commit run --all-files --verbose
       - name: Run tests 🧪

flashbax/buffers/prioritised_trajectory_buffer.py

@@ -798,6 +798,7 @@ def make_prioritised_trajectory_buffer(
     if max_size is not None:
         max_length_time_axis = max_size // add_batch_size
 
+    assert max_length_time_axis is not None
     init_fn = functools.partial(
         prioritised_init,
         add_batch_size=add_batch_size,

flashbax/buffers/trajectory_buffer.py

@@ -585,6 +585,7 @@ def make_trajectory_buffer(
     if max_size is not None:
         max_length_time_axis = max_size // add_batch_size
 
+    assert max_length_time_axis is not None
     init_fn = functools.partial(
         init,
         add_batch_size=add_batch_size,