Move AdvantageBuffer to its own file and switch to tensordict (#9)

* Changed:
- I'm going to switch to using tensordict as the base data structure in hopes
that it cleans up a lot of the API
- AdvantageBuffer is now separate

* Updated:
- firstlast merger now takes a merging function to support different
step merging operations. Might just rename the FirstLastMerger to a StepMerger,
but idk.
- API to use tensordict steps instead


* Added
- docker image build cleanup
- profile for the build section so that vscode doesnt
try rebuilding it when we can just build

* Updated:
- devcontainer stuff
- mojo stuff

* update docker since ubuntu 18 is deprecated

* Changed:
- quick updates

* Added:
- debug env for continuous agents

* Still skeptical whether advantage buffer works correctly

Author: josiahls

.devcontainer.json

@@ -2,18 +2,21 @@
     "name": "fastrl_development",
     "dockerComposeFile": "docker-compose.yml",
     "service": "fastrl",
-    // "settings": {"terminal.integrated.shell.linux": "/bin/bash",
-    //             "python.defaultInterpreterPath":"/opt/conda/bin/python"},
+    "customizations":{
+        "vscode": {
+            "settings": {"terminal.integrated.shell.linux": "/bin/bash"}
+        }
+    },
     // "mounts": [ "source=/var/run/docker.sock,target=/var/run/docker.sock,type=bind",
     //             "source=/usr/bin/docker,target=/usr/bin/docker,type=bind" ],
     "workspaceFolder": "/home/fastrl_user/fastrl", 
     // "forwardPorts": [4000, 8080],
     // "appPort": [4000, 8080],
-    "extensions": [
-		"ms-python.python",
-		"ms-azuretools.vscode-docker",
-		"ms-toolsai.jupyter-renderers"
-	],
+    // "extensions": [
+	// 	"ms-python.python",
+	// 	"ms-azuretools.vscode-docker",
+	// 	"ms-toolsai.jupyter-renderers"
+	// ],
     "runServices": ["dep_watcher", "quarto"] //,
     // "postStartCommand": "pip install -e .[dev]"
 }

.github/workflows/fastrl-docker.yml

@@ -7,7 +7,7 @@ on:
     branches:
       - main
       - update_nbdev_docs
-      - refactor/bug-fix-api-update-and-stablize
+      - refactor/advantage-buffer
 
 jobs:
   build:
@@ -18,10 +18,21 @@ jobs:
         build_type: [dev]
         # build_type: [prod, dev]
     steps:
+    - name: Maximize build space
+      uses: easimon/maximize-build-space@master
+      with:
+        root-reserve-mb: 35000
+        swap-size-mb: 1024
+        remove-dotnet: 'true'
+        remove-android: 'true'
     - name: Copy This Repository Contents
       uses: actions/checkout@v2
-      with:
-        submodules: recursive
+      # with:
+      #   submodules: recursive
+    - name: Build
+      run: |
+        echo "Free space:"
+        df -h
 
     - uses: actions/setup-python@v2
       with:
@@ -54,31 +65,36 @@ jobs:
       env:
         BUILD_TYPE: ${{ matrix.build_type }}
 
-    - name: Cache Docker layers
-      if: always()
-      uses: actions/cache@v2
-      with:
-        path: /tmp/.buildx-cache
-        key: ${{ runner.os }}-buildx-${{ github.sha }}
-        restore-keys: |
-          ${{ runner.os }}-buildx-
+    # - name: Cache Docker layers
+    #   if: always()
+    #   uses: actions/cache@v3
+    #   with:
+    #     path: /tmp/.buildx-cache
+    #     key: ${{ runner.os }}-buildx-${{ github.sha }}
+    #     restore-keys: |
+    #       ${{ runner.os }}-buildx-
 
-    - name: Set up Docker Buildx
-      uses: docker/setup-buildx-action@v3
+    # - name: Set up Docker Buildx
+    #   uses: docker/setup-buildx-action@v3
 
     - name: build and tag container
       run: |
         export DOCKER_BUILDKIT=1
         # We need to clear the previous docker images
-        # docker system prune -fa
-        # docker pull ${IMAGE_NAME}:latest || true
+        docker system prune -fa
+        docker pull ${IMAGE_NAME}:latest || true
         # docker build --build-arg BUILD=${BUILD_TYPE} \
-        docker buildx create --use
-        docker buildx build --cache-from=type=local,src=/tmp/.buildx-cache --cache-to=type=local,dest=/tmp/.buildx-cache --build-arg BUILD=${BUILD_TYPE} \
-          -t ${IMAGE_NAME}:latest \
-          -t ${IMAGE_NAME}:${VERSION} \
-          -t ${IMAGE_NAME}:$(date +%F) \
-          -f fastrl.Dockerfile .
+        # docker buildx create --use
+        # docker buildx build --cache-from=type=local,src=/tmp/.buildx-cache --cache-to=type=local,dest=/tmp/.buildx-cache --build-arg BUILD=${BUILD_TYPE} \
+        #   -t ${IMAGE_NAME}:latest \
+        #   -t ${IMAGE_NAME}:${VERSION} \
+        #   -t ${IMAGE_NAME}:$(date +%F) \
+        #   -f fastrl.Dockerfile .
+        docker buildx build --build-arg BUILD=${BUILD_TYPE} \
+              -t ${IMAGE_NAME}:latest \
+              -t ${IMAGE_NAME}:${VERSION} \
+              -t ${IMAGE_NAME}:$(date +%F) \
+              -f fastrl.Dockerfile .
       env:
         VERSION: ${{ steps.get_variables.outputs.version }}
         IMAGE_NAME: ${{ steps.get_variables.outputs.image_name }}

.github/workflows/fastrl-test.yml

@@ -1,33 +1,35 @@
 name: Fastrl Testing
 on: [push, pull_request]
 
-env:
-  PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/nightly/cu113
+# env:
+  # PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/nightly/cu113
 
 jobs:
-  build:
+  test:
     runs-on: ubuntu-latest
-    # container:
-    #   image: 'josiahls/fastrl-dev:latest'
+    container:
+      image: 'josiahls/fastrl-dev:latest'
 
     steps:
     - uses: actions/checkout@v1
-    - uses: actions/setup-python@v1
-      with:
-        python-version: '3.7'
-        architecture: 'x64'
+    # - uses: actions/setup-python@v1
+    #   with:
+    #     python-version: '3.7'
+        # architecture: 'x64'
     - name: Install the library
       run: |
-        pip install -e .["dev"]
+        sudo mkdir -p /github/home
+        sudo pip install -e .["dev"]
     # - name: Read all notebooks
     #   run: |
     #     nbdev_read_nbs
     - name: Check if all notebooks are cleaned
       run: |
+          sudo git config --global --add safe.directory /__w/fastrl/fastrl
           echo "Check we are starting with clean git checkout"
           if [ -n "$(git status -uno -s)" ]; then echo "git status is not clean"; false; fi
           echo "Trying to strip out notebooks"
-          nbdev_clean
+          sudo nbdev_clean
           echo "Check that strip out was unnecessary"
           git status -s # display the status to see which nbs need cleaning up
           if [ -n "$(git status -uno -s)" ]; then echo -e "!!! Detected unstripped out notebooks\n!!!Remember to run nbdev_install_hooks"; false; fi
@@ -37,7 +39,8 @@ jobs:
     #     if [ -n "$(nbdev_diff_nbs)" ]; then echo -e "!!! Detected difference between the notebooks and the library"; false; fi
     - name: Run tests
       run: |
-        pip show torchdata torch
+        pip3 show torchdata torch
+        sudo pip3 install -e .
         cd nbs 
         xvfb-run -s "-screen 0 1400x900x24" fastrl_nbdev_test --n_workers 12 --one2one
     - name: Run Doc Build Test

.github/workflows/quarto_deploy.yml

@@ -10,13 +10,11 @@ env:
 jobs:
   deploy:
     runs-on: ubuntu-latest
+    container:
+      image: 'josiahls/fastrl-dev:latest'
 
     steps: 
     - uses: actions/checkout@v3
-    - uses: actions/setup-python@v1
-      with:
-        python-version: '3.7'
-        architecture: 'x64'
     - name: Install Dependencies
       shell: bash
       run: |

docker-compose.yml

@@ -1,12 +1,16 @@
-version: "3"
+version: "3.8"
 services:
-  fastrl: &fastrl
-    restart: unless-stopped
-    working_dir: /home/fastrl_user/fastrl
+  fastrl_build:
     build:
       dockerfile: fastrl.Dockerfile
       context: .
     image: josiahls/fastrl-dev:latest
+    profiles: ["build"]
+
+  fastrl: &fastrl
+    restart: unless-stopped
+    working_dir: /home/fastrl_user/fastrl
+    image: josiahls/fastrl-dev:latest
     deploy:
       resources:
         reservations:

extra/pip_requirements.txt

@@ -1,4 +1,5 @@
 gymnasium>=0.27.1
+tensordict
 pyopengl
 pyglet
 tensorboard

fastrl.Dockerfile

@@ -1,6 +1,6 @@
 # FROM pytorch/pytorch:1.10.0-cuda11.3-cudnn8-runtime
 # FROM pytorch/pytorch:1.12.1-cuda11.3-cudnn8-runtime
-FROM nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu18.04
+FROM nvidia/cuda:12.0.0-runtime-ubuntu20.04
 # RUN  conda install python=3.8
 
 ENV CONTAINER_USER fastrl_user
@@ -12,13 +12,13 @@ RUN addgroup --gid $CONTAINER_UID $CONTAINER_GROUP && \
     #  && \
     # mkdir -p /opt/conda && chown $CONTAINER_USER /opt/conda
 
-RUN apt-get update && apt-get install -y software-properties-common rsync curl
+RUN apt-get update && apt-get install -y software-properties-common rsync curl gcc g++
 #RUN apt-key adv --keyserver keyserver.ubuntu.com --recv-key C99B11DEB97541F0 && apt-add-repository https://cli.github.com/packages
 
 RUN curl -fsSL https://cli.github.com/packages/githubcli-archive-keyring.gpg | gpg --dearmor -o /usr/share/keyrings/githubcli-archive-keyring.gpg
 RUN echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | tee /etc/apt/sources.list.d/github-cli.list > /dev/null
 
-RUN apt-get install -y python3.8-dev python3.8-distutils
+RUN apt-get install -y build-essential python3.8-dev python3.8-distutils
 RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.8 2 && update-alternatives --config python
 RUN curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3.8 get-pip.py
 
@@ -29,17 +29,21 @@ RUN apt-get update && apt-get install -y git libglib2.0-dev graphviz libxext6 \
 
 WORKDIR /home/$CONTAINER_USER
 # Install Primary Pip Reqs
-ENV PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/nightly/cu117
+# ENV PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/nightly/cu117
 COPY --chown=$CONTAINER_USER:$CONTAINER_GROUP extra/requirements.txt /home/$CONTAINER_USER/extra/requirements.txt
 # Since we are using a custom fork of torchdata, we install torchdata as part of a submodule.
 # RUN pip3 install -r extra/requirements.txt --pre --upgrade
-RUN pip3 install torch>=2.0.0 --pre --upgrade
+RUN pip3 install torch>=2.0.0 
+# --pre --upgrade
 RUN pip3 show torch
 
-
 COPY --chown=$CONTAINER_USER:$CONTAINER_GROUP extra/pip_requirements.txt /home/$CONTAINER_USER/extra/pip_requirements.txt
 RUN pip3 install -r extra/pip_requirements.txt
 
+WORKDIR /home/$CONTAINER_USER/fastrl
+RUN git clone https://github.com/josiahls/data.git \
+    && cd data && pip3 install -e .
+WORKDIR /home/$CONTAINER_USER
 
 # Install Dev Reqs
 COPY --chown=$CONTAINER_USER:$CONTAINER_GROUP extra/dev_requirements.txt /home/$CONTAINER_USER/extra/dev_requirements.txt
@@ -74,7 +78,7 @@ RUN /bin/bash -c "if [[ $BUILD == 'dev' ]] ; then nbdev_install_quarto ; fi"
 # ENV LD_LIBRARY_PATH /home/$CONTAINER_USER/.mujoco/mujoco210/bin:${LD_LIBRARY_PATH}
 # ENV LD_LIBRARY_PATH /usr/local/nvidia/lib64:${LD_LIBRARY_PATH}
 
-RUN ln -s /usr/lib/x86_64-linux-gnu/libGL.so.1 /usr/lib/x86_64-linux-gnu/libGL.so
+# RUN ln -s /usr/lib/x86_64-linux-gnu/libGL.so.1 /usr/lib/x86_64-linux-gnu/libGL.so
 
 USER $CONTAINER_USER
 WORKDIR /home/$CONTAINER_USER
@@ -84,9 +88,18 @@ ENV PATH="/home/$CONTAINER_USER/.local/bin:${PATH}"
 RUN pip install setuptools==60.7.0
 COPY --chown=$CONTAINER_USER:$CONTAINER_GROUP  . fastrl
 
-RUN sudo apt-get -y install cmake
+RUN sudo apt-get -y install cmake python3.8-venv
+
+# RUN curl https://get.modular.com | sh - && \
+#     modular auth mut_9b52dfea7b05427385fdeddc85dd3a64 && \
+#     modular install mojo
+
+RUN BASHRC=$( [ -f "$HOME/.bash_profile" ] && echo "$HOME/.bash_profile" || echo "$HOME/.bashrc" ) && \
+    echo 'export MODULAR_HOME="/home/fastrl_user/.modular"' >> "$BASHRC" && \
+    echo 'export PATH="/home/fastrl_user/.modular/pkg/packages.modular.com_mojo/bin:$PATH"' >> "$BASHRC" && \
+    source "$BASHRC"
 
-RUN /bin/bash -c "if [[ $BUILD == 'dev' ]] ; then echo \"Development Build\" && cd fastrl/data &&  mv pyproject.toml pyproject.toml_tmp && pip install -e . --no-dependencies &&  mv pyproject.toml_tmp pyproject.toml && cd ../; fi"
+# RUN /bin/bash -c "if [[ $BUILD == 'dev' ]] ; then echo \"Development Build\" && cd fastrl/data &&  mv pyproject.toml pyproject.toml_tmp && pip install -e . --no-dependencies &&  mv pyproject.toml_tmp pyproject.toml && cd ../; fi"
 
 RUN /bin/bash -c "if [[ $BUILD == 'prod' ]] ; then echo \"Production Build\" && cd fastrl && pip install . --no-dependencies; fi"
 RUN /bin/bash -c "if [[ $BUILD == 'dev' ]] ; then echo \"Development Build\" && cd fastrl && pip install -e \".[dev]\" --no-dependencies ; fi"