diff --git a/.gitignore b/.gitignore
index 19743fd0..1f366fb8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,6 @@
+# Environment file written by setup scripts
+pyrobosim.env
+
 # Dependencies folder
 dependencies/
 
diff --git a/docker/Dockerfile b/docker/Dockerfile
index b9351604..9f404b45 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -29,8 +29,8 @@ RUN mkdir -p /opt/pyrobosim/setup
 WORKDIR /opt/pyrobosim/
 
 # Install PDDLStream
-COPY setup/setup_pddlstream.bash setup/
-RUN setup/setup_pddlstream.bash
+COPY setup/configure_pddlstream.bash setup/
+RUN setup/configure_pddlstream.bash
 
 # Install pip dependencies for testing
 ENV PIP_BREAK_SYSTEM_PACKAGES=1
@@ -65,9 +65,9 @@ RUN apt-get install -y \
 RUN mkdir -p /pyrobosim_ws/src/pyrobosim
 
 # Install dependencies
-COPY setup/setup_pddlstream.bash /pyrobosim_ws/src/setup/
+COPY setup/configure_pddlstream.bash /pyrobosim_ws/src/setup/
 WORKDIR /pyrobosim_ws/src/setup
-RUN ./setup_pddlstream.bash
+RUN ./configure_pddlstream.bash
 
 COPY . /pyrobosim_ws/src/
 
diff --git a/docs/source/setup.rst b/docs/source/setup.rst
index 2e9420e0..dec1e678 100644
--- a/docs/source/setup.rst
+++ b/docs/source/setup.rst
@@ -26,18 +26,13 @@ To set up your Python virtual environment, configure and run
 
 ::
 
-  ./setup/create_python_env.bash
+  ./setup/setup_pyrobosim.bash
 
 By default, this will create a Python virtual environment in ``~/python-virtualenvs/pyrobosim``.
 
-If you want to use `PDDLStream <https://github.com/caelan/pddlstream>`_ for
-task and motion planning, you should also run:
+You will also get prompts for setting up ROS 2 and PDDLStream for task and motion planning.
 
-::
-
-   ./setup/setup_pddlstream.bash
-
-To then source this virtual environment, run
+To then setup the environment, run
 
 ::
 
@@ -51,28 +46,6 @@ As documented in the above script, we recommend making a bash function in your `
        source /path/to/pyrobosim/setup/source_pyrobosim.bash
     }
 
-Additional ROS 2 Setup
-----------------------
-
-After you have installed ``pyrobosim`` and activated your Python virtual environment,
-you must build your colcon workspace to install the ``pyrobosim_msgs`` and ``pyrobosim_ros`` packages.
-For example, if you have cloned this repo to ``~/pyrobosim_ws/src/pyrobosim``, you can do:
-
-::
-
-    cd ~/pyrobosim_ws
-    colcon build
-    . install/local_setup.bash
-
-For ROS 2 workflows, you can also make bash function to get set up like this:
-
-::
-
-    pyrobosim_ros() {
-       source /path/to/pyrobosim/setup/source_pyrobosim.bash humble
-    }
-
-The additional ``humble`` argument will make sure that ROS 2 Humble and your built colcon workspace are sourced in addition to activating the Python virtual environment.
 
 Docker Setup
 ------------
diff --git a/pyrobosim/examples/demo_astar.py b/pyrobosim/examples/demo_astar.py
index 83c80253..0e811d4e 100755
--- a/pyrobosim/examples/demo_astar.py
+++ b/pyrobosim/examples/demo_astar.py
@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+
 import os
 
 from pyrobosim.core import WorldYamlLoader
diff --git a/setup/setup_pddlstream.bash b/setup/configure_pddlstream.bash
similarity index 100%
rename from setup/setup_pddlstream.bash
rename to setup/configure_pddlstream.bash
diff --git a/setup/create_python_env.bash b/setup/create_python_env.bash
deleted file mode 100755
index 1f2b6dfd..00000000
--- a/setup/create_python_env.bash
+++ /dev/null
@@ -1,42 +0,0 @@
-#!/bin/bash
-
-# Sets up the pyrobosim virtual environment
-
-# User variables
-VIRTUALENV_FOLDER=~/python-virtualenvs/pyrobosim
-
-# Create a Python virtual environment
-[ ! -d "${VIRTUALENV_FOLDER}" ] && mkdir -p ${VIRTUALENV_FOLDER}
-python3 -m venv ${VIRTUALENV_FOLDER}
-echo "Created Python virtual environment in ${VIRTUALENV_FOLDER}"
-
-# Install all the Python packages required
-# Note that these overlay over whatever ROS 2 already contains
-source "${VIRTUALENV_FOLDER}/bin/activate"
-SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
-pushd "${SCRIPT_DIR}/.." > /dev/null || exit
-pip3 install ./pyrobosim
-pip3 install -r test/python_test_requirements.txt
-popd > /dev/null || exit
-deactivate
-
-# If setting up with ROS, install extra packages
-ROS_ARG=$1
-if [ "${ROS_ARG}" == "ros"]
-then
-  echo "Installing additional packages for ROS setup"
-  pip3 install colcon_common_extensions
-fi
-
-# Print confirmation and instructions at the end
-echo -e "
-
-
-Created Python virtual environment and installed packages at the following location:
-
-  ${VIRTUALENV_FOLDER}
-
-Source the environment using the following command:
-
-  source setup/source_pyrobosim.bash
-"
diff --git a/setup/setup_pyrobosim.bash b/setup/setup_pyrobosim.bash
new file mode 100755
index 00000000..16dde60a
--- /dev/null
+++ b/setup/setup_pyrobosim.bash
@@ -0,0 +1,69 @@
+#!/bin/bash
+
+# Sets up the pyrobosim virtual environment
+
+# User variables
+# Please modify these for your environment
+VIRTUALENV_FOLDER=~/python-virtualenvs/pyrobosim
+ROS_WORKSPACE=~/workspace/pyrobosim_ws
+
+# Create a Python virtual environment
+[ ! -d "${VIRTUALENV_FOLDER}" ] && mkdir -p ${VIRTUALENV_FOLDER}
+python3 -m venv ${VIRTUALENV_FOLDER}
+echo "Created Python virtual environment in ${VIRTUALENV_FOLDER}"
+
+# Install all the Python packages required
+# Note that these overlay over whatever ROS 2 already contains
+source "${VIRTUALENV_FOLDER}/bin/activate"
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+pushd "${SCRIPT_DIR}/.." > /dev/null
+pip3 install ./pyrobosim
+pip3 install -r test/python_test_requirements.txt
+
+# Write key variables to file
+ENV_FILE="pyrobosim.env"
+if [ -f "${ENV_FILE}" ]
+then
+  rm ${ENV_FILE}
+fi
+echo "# This is an autogenerated environment file for pyrobosim" >> ${ENV_FILE}
+echo "PYROBOSIM_VENV=${VIRTUALENV_FOLDER}" >> ${ENV_FILE}
+
+# If setting up with ROS, install extra packages
+echo -e ""
+read -p "Do you want to set up ROS? (y/n) : " USE_ROS
+if [ "${USE_ROS,,}" == "y" ]
+then
+  read -p "What ROS distro are you using? (humble, iron, jazzy, rolling) : " ROS_DISTRO
+  echo ""
+  echo "Installing additional packages for ROS ${ROS_DISTRO,,} setup"
+  echo "PYROBOSIM_ROS_WORKSPACE=${ROS_WORKSPACE}" >> ${ENV_FILE}
+  echo "PYROBOSIM_ROS_DISTRO=${ROS_DISTRO}" >> ${ENV_FILE}
+  pip3 install colcon_common_extensions
+fi
+
+# Optionally configure PDDLStream for task and motion planning
+echo -e ""
+read -p "Do you want to set up PDDLStream for task and motion planning? (y/n) : " USE_PDDLSTREAM
+if [ "${USE_PDDLSTREAM,,}" == "y" ]
+then
+  ./setup/configure_pddlstream.bash
+fi
+
+popd > /dev/null
+deactivate
+
+# Print confirmation and instructions at the end
+echo -e "
+
+
+Created Python virtual environment and installed packages at the following location:
+
+  ${VIRTUALENV_FOLDER}
+
+Environment data has been written to the file ${ENV_FILE}.
+
+Source the environment using the following command:
+
+  source setup/source_pyrobosim.bash
+"
diff --git a/setup/source_pyrobosim.bash b/setup/source_pyrobosim.bash
index 8713afed..a08d95ad 100755
--- a/setup/source_pyrobosim.bash
+++ b/setup/source_pyrobosim.bash
@@ -14,18 +14,21 @@
 #  So you can then run this from your Terminal:
 #    pyrobosim
 #
-# For ROS workflows, enter the ROS distro (humble/iron) as an argument:
-#
-#   pyrobosim_ros() {
-#     source /path/to/pyrobosim/setup/source_pyrobosim.bash humble
-#   }
-#
-#   So you can then call run from your Terminal:
-#     pyrobosim_ros
 
-# User variables -- change this to meet your needs
-export VIRTUALENV_FOLDER=~/python-virtualenvs/pyrobosim
-export PYROBOSIM_WS=~/workspace/pyrobosim_ws
+# Set up the environment
+ENV_FILE="pyrobosim.env"
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+pushd "${SCRIPT_DIR}/.." > /dev/null
+if [ ! -f "${ENV_FILE}" ]
+then
+    popd > /dev/null
+    echo "Did not find a file named ${ENV_FILE} in the root pyrobosim folder."
+    echo "Please rerun the setup_pyrobosim.bash script."
+    return 1
+fi
+unset PYROBOSIM_VENV PYROBOSIM_ROS_WORKSPACE PYROBOSIM_ROS_DISTRO
+source "${ENV_FILE}"
+popd > /dev/null
 
 if [ -n "${VIRTUAL_ENV}" ]
 then
@@ -33,30 +36,28 @@ then
 fi
 
 # Activate the Python virtual environment
-echo "Activated virtual environment at ${VIRTUALENV_FOLDER}."
-source ${VIRTUALENV_FOLDER}/bin/activate
+echo "Activated virtual environment at ${PYROBOSIM_VENV}."
+source ${PYROBOSIM_VENV}/bin/activate
 
 # Parse ROS distro argument
-ROS_DISTRO=$1
-if [ "${ROS_DISTRO}" == "" ]
+if [ -z "${PYROBOSIM_ROS_WORKSPACE}" ]
 then
     echo "Setting up pyrobosim with no ROS distro."
 else
-    echo "Setting up pyrobosim with ROS ${ROS_DISTRO}."
-    source /opt/ros/${ROS_DISTRO}/setup.bash
-    pushd "${PYROBOSIM_WS}" > /dev/null || exit
-    if [ ! -d "build" ]
+    echo "Setting up pyrobosim with ROS ${PYROBOSIM_ROS_DISTRO}."
+    source /opt/ros/${PYROBOSIM_ROS_DISTRO}/setup.bash
+    pushd "${PYROBOSIM_ROS_WORKSPACE}" > /dev/null
+    if [ ! -f "install/setup.bash" ]
     then
-        echo "Building ROS workspace at ${PYROBOSIM_WS}..."
+        echo "Building ROS workspace at ${PYROBOSIM_ROS_WORKSPACE}..."
         colcon build
     fi
-    echo "Sourcing ROS workspace at ${PYROBOSIM_WS}."
-    source install/local_setup.bash
-    popd > /dev/null || exit
+    echo "Sourcing ROS workspace at ${PYROBOSIM_ROS_WORKSPACE}."
+    source install/setup.bash
+    popd > /dev/null
 fi
 
 # Add dependencies to path
-SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 PDDLSTREAM_PATH=${SCRIPT_DIR}/../dependencies/pddlstream
 if [ -d "${PDDLSTREAM_PATH}" ]
 then