diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh
index 7ca5e70..42c6e2f 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh
@@ -1,13 +1,17 @@
-#!/bin/bash
-
-# Must exit and fail to build if any command fails
-set -eo pipefail
+#!/bin/false
 
 if [[ -z $ROCM_VERSION ]]; then
     printf "No valid ROCM_VERSION specified\n" >&2
     exit 1
 fi
 
+export ROCM_VERSION="$ROCM_VERSION"
+export ROCM_LEVEL="$ROCM_LEVEL"
+export PATH=/opt/rocm/bin:$PATH
+printf "export ROCM_VERSION=\"%s\"\n" "${ROCM_VERSION}" >> /opt/ai-dock/etc/environment.sh
+printf "export ROCM_LEVEL=\"%s\"\n" "${ROCM_LEVEL}" >> /opt/ai-dock/etc/environment.sh
+printf "export PATH=\"%s\"\n" "${PATH}" >> /etc/bash.bashrc
+
 curl -Ss https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor | tee /etc/apt/keyrings/rocm.gpg > /dev/null
 
 echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${ROCM_VERSION} jammy main" \
@@ -39,4 +43,4 @@ elif [[ "${ROCM_LEVEL}" == "devel" ]]; then
 else
     printf "No valid ROCM_LEVEL specified\n" >&2
     exit 1
-fi
\ No newline at end of file
+fi
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/clean.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/clean.sh
index 0e167bd..ef37d74 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/clean.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/clean.sh
@@ -1,7 +1,5 @@
-#!/bin/bash
+#!/bin/false
 
-# Must exit and fail to build if any command fails
-set -eo pipefail
 
 # Tidy up and keep image small
 apt-get clean -y
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/common.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/common.sh
index 4186a79..95bf8da 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/common.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/common.sh
@@ -1,7 +1,9 @@
-#!/bin/bash
+#!/bin/false
 
-# Must exit and fail to build if any command fails
-set -eo pipefail
+export MAMBA_CREATE="micromamba create --always-softlink -y -c conda-forge"
+export MAMBA_INSTALL="micromamba install --always-softlink -y -c conda-forge"
+printf "export MAMBA_CREATE=\"%s\"\n" "${MAMBA_CREATE}" >> /opt/ai-dock/etc/environment.sh
+printf "export MAMBA_INSTALL=\"%s\"\n" "${MAMBA_INSTALL}" >> /opt/ai-dock/etc/environment.sh
 
 apt-get update
 apt-get upgrade -y --no-install-recommends
@@ -92,3 +94,8 @@ touch /etc/rclone/rclone.conf
 # Git config
 
 git config --global --add safe.directory "*"
+
+# Ensure correct environment for child builds
+
+printf "source /opt/ai-dock/etc/environment.sh\n" >> /etc/profile.d/02-ai-dock.sh
+printf "source /opt/ai-dock/etc/environment.sh\n" >> /etc/bash.bashrc
\ No newline at end of file
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/config.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/config.sh
deleted file mode 100644
index 21cdf1b..0000000
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/config.sh
+++ /dev/null
@@ -1,3 +0,0 @@
-#!/bin/false
-
-# A tidy place to set some variables
\ No newline at end of file
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/cpu.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/cpu.sh
index 4fbd580..8e812c3 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/cpu.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/cpu.sh
@@ -1,6 +1,3 @@
-#!/bin/bash
+#!/bin/false
 
-# Must exit and fail to build if any command fails
-set -eo pipefail
 
-exit 0
\ No newline at end of file
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/init.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/init.sh
index 9c80db1..d943313 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/init.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/init.sh
@@ -3,19 +3,17 @@
 # Must exit and fail to build if any command fails
 set -eo pipefail
 
-source /opt/ai-dock/bin/build/layer0/config.sh
-
-/opt/ai-dock/bin/build/layer0/common.sh
+source /opt/ai-dock/bin/build/layer0/common.sh
 
 if [[ "$XPU_TARGET" == "NVIDIA_GPU" ]]; then
-    /opt/ai-dock/bin/build/layer0/nvidia.sh
+    source /opt/ai-dock/bin/build/layer0/nvidia.sh
 elif [[ "$XPU_TARGET" == "AMD_GPU" ]]; then
-    /opt/ai-dock/bin/build/layer0/amd.sh
+    source /opt/ai-dock/bin/build/layer0/amd.sh
 elif [[ "$XPU_TARGET" == "CPU" ]]; then
-    /opt/ai-dock/bin/build/layer0/cpu.sh
+    source /opt/ai-dock/bin/build/layer0/cpu.sh
 else
     printf "No valid XPU_TARGET specified\n" >&2
     exit 1
 fi
 
-/opt/ai-dock/bin/build/layer0/clean.sh
\ No newline at end of file
+source /opt/ai-dock/bin/build/layer0/clean.sh
\ No newline at end of file
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh
index 4fbd580..6596b6a 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh
@@ -1,6 +1,14 @@
-#!/bin/bash
+#!/bin/false
 
-# Must exit and fail to build if any command fails
-set -eo pipefail
-
-exit 0
\ No newline at end of file
+export CUDA_VERSION=${CUDA_VERSION}
+export CUDNN_VERSION=${CUDNN_VERSION}
+export CUDA_LEVEL=${CUDA_LEVEL}
+export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
+export MAMBA_CREATE="micromamba create --always-softlink -y -c nvidia -c conda-forge"
+export MAMBA_INSTALL="micromamba install --always-softlink -y -c nvidia -c conda-forge"
+printf "export CUDA_VERSION=\"%s\"\n" "${CUDA_VERSION}" >> /opt/ai-dock/etc/environment.sh
+printf "export CUDNN_VERSION=\"%s\"\n" "${CUDNN_VERSION}" >> /opt/ai-dock/etc/environment.sh
+printf "export CUDA_LEVEL=\"%s\"\n" "${CUDA_LEVEL}" >> /opt/ai-dock/etc/environment.sh
+printf "export LD_LIBRARY_PATH=\"%s\"\n" "${LD_LIBRARY_PATH}" >> /opt/ai-dock/etc/environment.sh
+printf "export MAMBA_CREATE=\"%s\"\n" "${MAMBA_CREATE}" >> /opt/ai-dock/etc/environment.sh
+printf "export MAMBA_INSTALL=\"%s\"\n" "${MAMBA_INSTALL}" >> /opt/ai-dock/etc/environment.sh
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/init.sh b/build/COPY_ROOT/opt/ai-dock/bin/init.sh
index e42726a..bf84c1b 100755
--- a/build/COPY_ROOT/opt/ai-dock/bin/init.sh
+++ b/build/COPY_ROOT/opt/ai-dock/bin/init.sh
@@ -84,11 +84,6 @@ function init_set_envs() {
     while IFS='=' read -r -d '' key val; do
         export "${key}"="$(init_strip_quotes "${val//___/' '}")"
     done < <(env -0)
-    
-    # TODO: branch init.sh into common,nvidia,amd,cpu
-    if [[ $XPU_TARGET == "AMD_GPU" ]]; then
-        export PATH=$PATH:/opt/rocm/bin
-    fi
 }
 
 function init_set_ssh_keys() {
@@ -396,11 +391,10 @@ function init_source_preflight_script() {
 
 function init_write_environment() {
     # Ensure all variables available for interactive sessions
-    env > /etc/environment
     while IFS='=' read -r -d '' key val; do
-        printf "export %s=\"%s\"\n" "$key" "$val" >> /root/.bashrc
+        printf "export %s=\"%s\"\n" "$key" "$val" >> /opt/ai-dock/etc/environment.sh
     done < <(env -0)
-    
+
     if [[ -n $MAMBA_DEFAULT_ENV ]]; then
       printf "micromamba activate %s\n" $MAMBA_DEFAULT_ENV >> /root/.bashrc
     fi
@@ -460,6 +454,9 @@ function init_debug_print() {
         printf "authorized_keys...\n\n"
         cat /root/.ssh/authorized_keys
         printf "\n--------------------------------------------\n"
+        printf "/opt/ai-dock/etc/environment.sh...\n\n"
+        cat /opt/ai-dock/etc/environment.sh
+        printf "\n--------------------------------------------\n"
         printf ".bashrc...\n\n"
         cat /root/.bashrc
         printf "\n---------- END DEBUG INFO---------- \n\n\n"
diff --git a/build/COPY_ROOT/opt/ai-dock/etc/environment.sh b/build/COPY_ROOT/opt/ai-dock/etc/environment.sh
new file mode 100644
index 0000000..e69de29
diff --git a/build/Dockerfile b/build/Dockerfile
index 92b54cd..6083978 100644
--- a/build/Dockerfile
+++ b/build/Dockerfile
@@ -18,33 +18,23 @@ SHELL ["/bin/bash", "-c"]
 ENV SHELL="/bin/bash"
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PATH=/opt/ai-dock/bin:/opt/micromamba/bin:/opt/caddy/bin:$PATH
-ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
 ENV OPT_SYNC=
 ENV PYTHONUNBUFFERED=true
 ENV APT_INSTALL="apt-get install -y --no-install-recommends"
 ENV PIP_INSTALL="pip install --no-cache-dir"
-ENV MAMBA_CHANNELS="-c nvidia -c conda-forge"
-ENV MAMBA_CREATE="micromamba create --always-softlink -y ${MAMBA_CHANNELS}"
-ENV MAMBA_INSTALL="micromamba install --always-softlink -y ${MAMBA_CHANNELS}"
 ENV MAMBA_ROOT_PREFIX=/opt/micromamba
 ENV MAMBA_EXE="/opt/micromamba/bin/micromamba"
 ENV RCLONE_CONFIG="/etc/rclone/rclone.conf"
+ENV IMAGE_SLUG="base-image"
 # Copy early so we can use scripts in the build - Changes to these files will invalidate the cache and cause a rebuild.
 COPY ./COPY_ROOT/ /
 
-# Env these args - We need them to ease software installation in extended images
 ARG CUDA_VERSION
-ENV CUDA_VERSION=${CUDA_VERSION}
 ARG CUDNN_VERSION
-ENV CUDNN_VERSION=${CUDNN_VERSION}
 ARG CUDA_LEVEL
-ENV CUDA_LEVEL=${CUDA_LEVEL}
 ARG ROCM_VERSION
-ENV ROCM_VERSION=${ROCM_VERSION}
 ARG ROCM_LEVEL
-ENV ROCM_LEVEL=${ROCM_LEVEL}
 
-ENV IMAGE_SLUG="base-image"
 # Use build scripts to ensure we can build all targets from one Dockerfile in a single layer.
 # Don't put anything heavy in here - We can use multi-stage building above if necessary.
 RUN set -eo pipefail && /opt/ai-dock/bin/build/layer0/init.sh | tee /var/log/build.log