diff --git a/docker/Dockerfile b/docker/Dockerfile
index f5102a4..443c8dd 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,12 +1,20 @@
-FROM  nvcr.io/nvidia/nvhpc:26.5-devel-cuda13.2-ubuntu24.04
+FROM nvcr.io/nvidia/nvhpc:26.5-devel-cuda13.2-ubuntu24.04 AS base
+
+FROM base AS arch-amd64
+ENV ARCH=x86_64
+
+FROM base AS arch-arm64
+ENV ARCH=aarch64
+
+FROM arch-${TARGETARCH} AS final
 
 SHELL ["/bin/bash", "-c"]
 
-ENV CUDA_HOME /opt/nvidia/hpc_sdk/Linux_x86_64/26.5/cuda
-ENV LD_LIBRARY_PATH /opt/nvidia/hpc_sdk/Linux_x86_64/26.5/cuda/13.2/extras/CUPTI/lib64:$LD_LIBRARY_PATH
+ENV CUDA_HOME=/opt/nvidia/hpc_sdk/Linux_${ARCH}/26.5/cuda
+ENV LD_LIBRARY_PATH=/opt/nvidia/hpc_sdk/Linux_${ARCH}/26.5/cuda/13.2/extras/CUPTI/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
 
 # Install System Dependencies
-ENV DEBIAN_FRONTEND noninteractive
+ENV DEBIAN_FRONTEND=noninteractive
 RUN apt update -y && \
     apt install -y wget cmake && \
     apt install -y python3 python-is-python3 python3-pip python3-pybind11 && \
@@ -18,10 +26,10 @@ RUN pip3 install --break-system-packages typing_extensions && \
     pip3 install --break-system-packages --no-deps nvidia-cudnn-cu13==9.20.0.48 nvidia-cusparselt-cu13==0.8.1 nvidia-nccl-cu13==2.29.7 nvidia-cufile==1.18.1.6
 
 # Remove conflicting NCCL version from NVHPC SDK, add libnccl.so symlink to pip installed NCCL
-RUN rm -rf /opt/nvidia/hpc_sdk/Linux_x86_64/26.5/comm_libs/nccl/ && \
+RUN rm -rf /opt/nvidia/hpc_sdk/Linux_${ARCH}/26.5/comm_libs/nccl/ && \
     cd /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib && \
     ln -s libnccl.so.2 libnccl.so
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH
+ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH
 
 # Install yaml-cpp
 RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
@@ -32,7 +40,7 @@ RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
           -DBUILD_SHARED_LIBS=OFF \
           -DCMAKE_POSITION_INDEPENDENT_CODE=ON .. && \
     make -j$(nproc) && make install
-ENV LD_LIBRARY_PATH /opt/yaml-cpp/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH=/opt/yaml-cpp/lib:${LD_LIBRARY_PATH}
 
 # Install additional Python dependencies
 RUN pip3 install --break-system-packages wandb ruamel-yaml matplotlib pygame moviepy
@@ -52,5 +60,5 @@ RUN cd /torchfort && mkdir build && cd build && \
     .. && \
     make -j$(nproc) install && \
     cd / && rm -rf torchfort
-ENV LD_LIBRARY_PATH /opt/torchfort/lib:${LD_LIBRARY_PATH}
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH=/opt/torchfort/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
diff --git a/docker/Dockerfile_gnu b/docker/Dockerfile_gnu
index a40bdee..9894342 100644
--- a/docker/Dockerfile_gnu
+++ b/docker/Dockerfile_gnu
@@ -1,11 +1,19 @@
-FROM nvcr.io/nvidia/cuda:13.2.0-devel-ubuntu24.04
+FROM nvcr.io/nvidia/cuda:13.2.0-devel-ubuntu24.04 AS base
+
+FROM base AS arch-amd64
+ENV ARCH=x86_64
+
+FROM base AS arch-arm64
+ENV ARCH=aarch64
+
+FROM arch-${TARGETARCH} AS final
 
 SHELL ["/bin/bash", "-c"]
 
-ENV CUDA_HOME /usr/local/cuda
+ENV CUDA_HOME=/usr/local/cuda
 
 # Install System Dependencies
-ENV DEBIAN_FRONTEND noninteractive
+ENV DEBIAN_FRONTEND=noninteractive
 RUN apt update -y && \
     apt install -y wget cmake && \
     apt install -y python3 python-is-python3 python3-pip python3-pybind11 && \
@@ -14,13 +22,13 @@ RUN apt update -y && \
 
 # Download HPCX
 RUN cd /opt && \
-    wget https://content.mellanox.com/hpc/hpc-x/v2.24.1_cuda13/hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64.tbz && \
-    tar xjf hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64.tbz && \
-    mv hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64 hpcx && \
-    cd /opt && rm hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64.tbz
+    wget https://content.mellanox.com/hpc/hpc-x/v2.24.1_cuda13/hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-${ARCH}.tbz && \
+    tar xjf hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-${ARCH}.tbz && \
+    mv hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-${ARCH} hpcx && \
+    cd /opt && rm hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-${ARCH}.tbz
 
-ENV PATH /opt/hpcx/ompi/bin:$PATH
-ENV LD_LIBRARY_PATH /opt/hpcx/ompi/lib:$LD_LIBRARY_PATH
+ENV PATH=/opt/hpcx/ompi/bin:$PATH
+ENV LD_LIBRARY_PATH=/opt/hpcx/ompi/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
 
 RUN echo "source /opt/hpcx/hpcx-init.sh; hpcx_load" >>  /root/.bashrc
 
@@ -32,9 +40,9 @@ RUN pip3 install --break-system-packages typing_extensions && \
 # Add libnccl.so symlink to pip installed NCCL (CUDA devel image ships no NCCL)
 RUN cd /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib && \
     ln -s libnccl.so.2 libnccl.so
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH
+ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH
 # NVSHMEM is a torch runtime dependency (libnvshmem_host.so.3) not present in the CUDA image
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/nvidia/nvshmem/lib:$LD_LIBRARY_PATH
+ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/nvidia/nvshmem/lib:$LD_LIBRARY_PATH
 
 # Install yaml-cpp
 RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
@@ -45,7 +53,7 @@ RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
           -DBUILD_SHARED_LIBS=OFF \
           -DCMAKE_POSITION_INDEPENDENT_CODE=ON .. && \
     make -j$(nproc) && make install
-ENV LD_LIBRARY_PATH /opt/yaml-cpp/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH=/opt/yaml-cpp/lib:${LD_LIBRARY_PATH}
 
 # Install additional Python dependencies
 RUN pip3 install --break-system-packages wandb ruamel-yaml matplotlib pygame moviepy
@@ -64,5 +72,5 @@ RUN source /opt/hpcx/hpcx-init.sh &&  hpcx_load && \
     .. && \
     make -j$(nproc) install && \
     cd / && rm -rf torchfort
-ENV LD_LIBRARY_PATH /opt/torchfort/lib:${LD_LIBRARY_PATH}
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH=/opt/torchfort/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
diff --git a/docker/Dockerfile_gnu_cpuonly b/docker/Dockerfile_gnu_cpuonly
index 29a3bab..83c4dc2 100644
--- a/docker/Dockerfile_gnu_cpuonly
+++ b/docker/Dockerfile_gnu_cpuonly
@@ -3,7 +3,7 @@ FROM ubuntu:24.04
 SHELL ["/bin/bash", "-c"]
 
 # Install System Dependencies
-ENV DEBIAN_FRONTEND noninteractive
+ENV DEBIAN_FRONTEND=noninteractive
 RUN apt update -y && \
     apt install -y build-essential && \
     apt install -y wget cmake && \
@@ -24,7 +24,7 @@ RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
           -DBUILD_SHARED_LIBS=OFF \
           -DCMAKE_POSITION_INDEPENDENT_CODE=ON .. && \
     make -j$(nproc) && make install
-ENV LD_LIBRARY_PATH /opt/yaml-cpp/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH=/opt/yaml-cpp/lib
 
 # Install additional Python dependencies
 RUN pip3 install --break-system-packages wandb ruamel-yaml matplotlib pygame moviepy
@@ -42,5 +42,5 @@ RUN cd /torchfort && mkdir build && cd build && \
     .. && \
     make -j$(nproc) install && \
     cd / && rm -rf torchfort
-ENV LD_LIBRARY_PATH /opt/torchfort/lib:${LD_LIBRARY_PATH}
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH=/opt/torchfort/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
diff --git a/tests/rl/test_replay_buffer.cpp b/tests/rl/test_replay_buffer.cpp
index f7da601..2296c20 100644
--- a/tests/rl/test_replay_buffer.cpp
+++ b/tests/rl/test_replay_buffer.cpp
@@ -258,7 +258,14 @@ TEST_P(ReplayBuffer, NStepConsistency) {
   }
 
   EXPECT_FLOAT_EQ(sdiff, 0.);
-  EXPECT_FLOAT_EQ(rdiff, 0.);
+  // rdiff accumulates float32 rounding from the n-step discounted reward sum
+  // (r0 + r1*gamma + r2*gamma^2 + ...) over the whole batch. The buffer and this
+  // manual recomputation use mathematically equivalent but differently-rounded
+  // paths (std::pow + torch tensor ops vs. repeated scalar float multiply), which
+  // diverge by a few ULP per term and differ across architectures (e.g. FMA
+  // contraction on aarch64). Compare against a small absolute tolerance rather
+  // than exact float equality.
+  EXPECT_NEAR(rdiff, 0., 1e-4);
 }
 
 TEST_P(ReplayBuffer, SaveRestore) {