Repository: buildingamind/NewbornEmbodiedTuringTest (language: python)

Issue:
Add Link to Existing Executables
Add a link to the URL containing existing executables from the lab to the README and to the documentation.


Discussion / hints:

Fix patch:
diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
index dd84ea78..94ddb904 100644
--- a/.github/ISSUE_TEMPLATE/bug_report.md
+++ b/.github/ISSUE_TEMPLATE/bug_report.md
@@ -25,13 +25,6 @@ If applicable, add screenshots to help explain your problem.
 
 **Desktop (please complete the following information):**
  - OS: [e.g. iOS]
- - Browser [e.g. chrome, safari]
- - Version [e.g. 22]
-
-**Smartphone (please complete the following information):**
- - Device: [e.g. iPhone6]
- - OS: [e.g. iOS8.1]
- - Browser [e.g. stock browser, safari]
  - Version [e.g. 22]
 
 **Additional context**
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
deleted file mode 100644
index 6894b2c2..00000000
--- a/.github/workflows/ci.yml
+++ /dev/null
@@ -1,25 +0,0 @@
-name: ci
-on:
-  push:
-    branches:
-      - main
-      - dev
-permissions:
-  contents: write
-jobs:
-  deploy:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v4
-        with:
-          python-version: 3.x
-      - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV
-      - uses: actions/cache@v3
-        with:
-          key: mkdocs-material-${{ env.cache_id }}
-          path: .cache
-          restore-keys: |
-            mkdocs-material-
-      - run: pip install mkdocs-material
-      - run: mkdocs gh-deploy --force   
\ No newline at end of file
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
new file mode 100644
index 00000000..110631fe
--- /dev/null
+++ b/.github/workflows/docs.yml
@@ -0,0 +1,55 @@
+name: "Render Docs"
+
+on:
+  push:
+    branches:
+      - main
+      - dev
+  pull_request:
+    branches:
+      - main
+      - dev
+
+jobs:
+  docs:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up Python 3.10.12
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.10.12
+
+      - name: Cache pip
+        uses: actions/cache@v3
+        id: cache-pip
+        with:
+          path: ${{ env.pythonLocation }}
+          key: ${{ runner.os }}-pip-${{ hashFiles('docs/requirements.txt') }}
+          restore-keys: |
+            ${{ runner.os }}-pip
+      - name: Install pip
+        if: steps.cache-pip.outputs.cache-hit != 'true'
+        run: |
+          pip install setuptools==65.5.0 pip==21
+      - name: Install Dependencies
+        if: steps.cache-pip.outputs.cache-hit != 'true'
+        run: |
+          pip install -r docs/requirements.txt
+
+      - name: Sphinx build
+        run: |
+          sphinx-build --jobs auto -b html docs/source/ docs/build/
+
+      - name: Deploy to GitHub Pages
+        uses: peaceiris/actions-gh-pages@v3
+        # if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }}
+        with:
+          publish_branch: gh-pages
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          publish_dir: docs/build/
+          force_orphan: true
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
new file mode 100644
index 00000000..2a0ae476
--- /dev/null
+++ b/.github/workflows/lint.yml
@@ -0,0 +1,47 @@
+name: Lint
+
+on:
+  push:
+    branches:
+      - main
+      - dev
+  pull_request:
+    branches:
+      - main
+      - dev
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10.12"
+          cache: "pip" # caching pip dependencies
+            # Install your linters here
+
+      # - name: Change setuptools version
+      #   run: pip install setuptools==65.5.0
+
+      # - name: Change pip version
+      #   run: pip install pip==21
+
+      - name: Install pylint
+        run: pip install pylint
+
+      # - name: Install dependencies
+      #   run: pip install -r requirements.txt
+
+      - name: Run Linter
+        uses: wearerequired/lint-action@v2
+        with:
+          pylint: true
+          pylint_args: "--disable=all --enable=C0114 --enable=C0115 --enable=C0116 --enable=W0612 --enable=W0613"
+          # checks for:
+          ## missing module docstring, missing class docstring, missing function docstring,
+          ## unused arguments, and unused variables
diff --git a/.gitignore b/.gitignore
index 36e1a652..8a87b1b3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -157,3 +157,57 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+
+# R code
+
+# History files
+.Rhistory
+.Rapp.history
+
+# Session Data files
+.RData
+.RDataTmp
+
+# User-specific files
+.Ruserdata
+
+# Example code in package build process
+*-Ex.R
+
+# Output files from R CMD build
+/*.tar.gz
+
+# Output files from R CMD check
+/*.Rcheck/
+
+# RStudio files
+.Rproj.user/
+
+# produced vignettes
+vignettes/*.html
+vignettes/*.pdf
+
+# OAuth2 token, see https://github.com/hadley/httr/releases/tag/v0.3
+.httr-oauth
+
+# knitr and R markdown default cache directories
+*_cache/
+/cache/
+
+# Temporary files created by R markdown
+*.utf8.md
+*.knit.md
+
+# R Environment Variables
+.Renviron
+
+# translation temp files
+po/*~
+
+# RStudio Connect folder
+rsconnect/
+
+# MacOS .DS_Store files
+.DS_Store
+# Remove .vscode folder
+.vscode/
diff --git a/LICENSE b/LICENSE
index 27417921..03085e9f 100644
--- a/LICENSE
+++ b/LICENSE
@@ -2,7 +2,6 @@
 MIT License
 
 Copyright (c) 2023 Building A Mind Lab
->>>>>>> 8cdfb7978c954597cb94be383b38228f3e6f342e
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 00000000..b97e47c0
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,3 @@
+# needed to include non-python files to installation
+include src/nett/analysis/*.R
+include src/nett/analysis/ChickData/*.csv
diff --git a/README.md b/README.md
index 7cba1467..b805fee5 100644
--- a/README.md
+++ b/README.md
@@ -1,138 +1,155 @@
-# **Unity environment for ChickAI: virtual controlled-rearing experiments**
+<div align="center">
+<img src="docs/assets/images/banner.png" alt="Banner" style />
 
-This is a collection of tools for simulating virtual agents under controlled-rearing conditions. The agents
-generated and studied through this pipeline can be compared directly to real chicks recorded by the **[**Building a Mind
-Lab**](http://buildingamind.com/)**. This pipeline provides all necessary components for simulating and replicating embodied models from the lab.
+# **Newborn Embodied Turing Test**
 
-The figure below shows the experiment setup for the three experiments discussed in the guide.
+Benchmarking Virtual Agents in Controlled-Rearing Conditions
 
-<img src="docs/digital_twin.jpg" alt="Digital Twin" style="zoom:35%;" />
+![PyPI - Version](https://img.shields.io/pypi/v/nett-benchmarks)
+![Python Version from PEP 621 TOML](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2Fbuildingamind%2FNewbornEmbodiedTuringTest%2Fmain%2Fpyproject.toml)
+![GitHub Issues or Pull Requests](https://img.shields.io/github/issues/buildingamind/NewbornEmbodiedTuringTest)
 
-## **How to Use this Repository**
+[Getting Started](#getting-started) •
+[Documentation](https://buildingamind.github.io/NewbornEmbodiedTuringTest/) • 
+[Lab Website](http://buildingamind.com/)
 
-This directory provides three components for building embodied virtual agents. These are a video game which serves as a virtual world, a set of programs to run experiments in the virtual world, and a set of programs to visualize the data coming from the experiments.
+</div>
 
-## **Directory Structure**
+The Newborn Embodied Turing Test (NETT) is a cutting-edge toolkit designed to simulate virtual agents in controlled-rearing conditions. This innovative platform enables researchers to create, simulate, and analyze virtual agents, facilitating direct comparisons with real chicks as documented by the **[Building a Mind Lab](http://buildingamind.com/)**. Our comprehensive suite includes all necessary components for the simulation and analysis of embodied models, closely replicating laboratory conditions.
 
-**Following the directory structure of the code.**
+Below is a visual representation of our experimental setup, showcasing the infrastructure for the three primary experiments discussed in this documentation.
 
-```
-
-├── docs
-├── mkdocs.yml
-├── README.md
-├── notebooks
-│   ├── Getting Started.ipynb
-├── src/netts
-│   ├── body
-│   └── brain
-│   ├── environment
-│   └── utils
-│   ├── nett.py
-│   └── __init__.py
-└── tests
-
-```
-
-* `src/netts`**: **Contains the code for running experiments with simulated agents. Following is the structure of `src/netts` folder:
-* `tests`: Contains unit tests
-* `docs`: Contains project documentation 
+<div align="center">
+<img src="docs/assets/images/digital_twin_cropped.jpg" alt="Digital Twin" width="65%" />
+</div>
 
-## **Getting Started**
-
-In this section, you will learn to use the repository to benchmark your first embodied agent with NETTs! 
-
-### **Codebase Installation**
-
-1. **(Highly Recommended) **[create and configure a virtual environment](https://uoa-eresearch.github.io/eresearch-cookbook/recipe/2014/11/20/conda/ "Link for how to set-up a virtual env")
-   ****steps described below :****
-   ```
-   conda create -n netts python=3.10.12
-   conda activate netts
-   pip install setuptools==65.5.0 pip==21
-   ```
-2. To install the repository using `pip`:
-   ```
-   pip install nett-benchmarks
-   ```
+## How to Use this Repository
 
-Note if not installing in a virtual environment, the install might fail because of conflicting dependency versions. `mlagents` uses `gym==0.21` which in-turn has dependencies on older versions of `numpy`, specifically `1.21.2` or below. Ensure that these requirements are met in your environment before proceeding. 
+The NETT toolkit comprises three key components:
 
-### **Running a NETT**
+1. **Virtual Environment**: A dynamic environment that serves as the habitat for virtual agents.
+2. **Experimental Simulation Programs**: Tools to initiate and conduct experiments within the virtual world.
+3. **Data Visualization Programs**: Utilities for analyzing and visualizing experiment outcomes.
 
-After having followed steps above, NETTs can be run with a few lines of code:
+## Directory Structure
 
-1. Download the executable from the hosted webpage or build your own executable following the steps mentioned in {placeholder}. Additionally, a video walkthrough of creating an executable from scratch is made available {here}. 
+The directory structure of the code is as follows:
 
-2. The package defines three components–the `Brain`, `Body` and `Environment`. For a detailed description of the division of responsibilities of each component, please refer to the {documentation}. 
 ```
-from netts import Brain, Body, Environment
-from netts import NETT
+├── docs                       # Documentation and guides
+├── notebooks
+│   ├── Getting Started.ipynb  # Introduction and setup notebook
+├── src/nett
+│   ├── analysis               # Analysis scripts
+│   ├── body                   # Agent body configurations
+│   ├── brain                  # Neural network models and learning algorithms
+│   ├── environment            # Simulation environments
+│   ├── utils                  # Utility functions
+│   ├── nett.py                # Main library script
+│   └── __init__.py            # Package initialization
+├── tests                      # Unit tests
+├── mkdocs.yml                 # MkDocs configuration
+├── pyproject.toml             # Project metadata
+└── README.md                  # This README file
 ```
 
-3. Define each component as required. Let’s start with the `Brain`. This component holds all the parts concerned with “learning”. This includes the architecture itself along with its many parameters, as well as the reward function and learning algorithm, such as `PPO`. The package is designed to make each component flexible. Specifically, each constituent of the Brain such as the encoder, policy networks or the reward function can be customized. The training and testing loops can be customized by inheriting the class and overriding them. This may be necessary in specialized cases, such as running on customized hardware such as TPU and IPUs. More details can be found in the documentation. 
+## Getting Started
 
-Consider a rather simple definition: 
-```
-brain = Brain(policy="CnnPolicy", algorithm="PPO")
-```
+To begin benchmarking your first embodied agent with NETT, please be aware:
 
-This defines the policy network and the learning algorithm, the rest of the arguments are left to defaults. Please refer to the repository for a complete list of supported arguments and customizations. This document will be updated with appropriate links to tutorials planned as part of future releases too. Note that the repository uses modules from `stable-baselines3` underneath and the value specified as arguments directly correspond with it. 
+**Important**: The `mlagents==1.0.0` dependency is incompatible with Apple Silicon (M1, M2, etc.) chips. Please utilize an alternate device to execute this codebase.
 
-4. Next, define the `Body`. Technically, the body is a medium through which all information passes through before reaching the brain. Hence, this component is primarily concerned with the application of `gym.Wrappers` (such as the DVS wrapper) that modify the information from the environment before they “reach” the brain for processing. 
+### Installation
 
-Consider a simple definition: 
-```
-body = Body(type="basic", dvs=False, wrappers=None)
-```
+1. **Virtual Environment Setup** (Highly Recommended): Create and activate a virtual environment to avoid dependency conflicts.
+   ```bash
+   conda create -y -n nett_env python=3.10.12
+   conda activate nett_env
+   ```
+   See [here](https://uoa-eresearch.github.io/eresearch-cookbook/recipe/2014/11/20/conda "Link for how to set-up a virtual env") for detailed instructions.
 
-In this example, we do not pass any wrappers. Or, we let the information from the environment reach the brain "as is". 
+2. Install the needed versions of `setuptools` and `pip`:
+   ```bash
+   pip install setuptools==65.5.0 pip==21 wheel==0.38.4
+   ```
+   **NOTE:** This is a result of incompatibilities with the subdependency `gym==0.21`. More information about this issue can be found [here](https://github.com/openai/gym/issues/3176#issuecomment-1560026649)
 
-While the `Body` abstraction is thin and mostly conceptual as of now, different types of agents (`two-eyed`, `rag-doll`) with complex action spaces are planned which lead to `Body` taking on more sophisticated customization. 
+3. **Toolkit Installation**: Install the toolkit using `pip`.
+   ```bash
+   pip install nett-benchmarks
+   ```
 
-5. The `Environment` component constructs a live environment (via the Python `mlagents` library) which is then wrapped inside a `Gym` environment. Since this is part of the environment initialization, the wrapping is not included as part of the `Body` component, but kept within the `Environment`. The definition takes an executable path which must be available on the system. 
-```
-# a valid Unity executable downloaded from {} must be present at `executable_path`
-environment = Environment(config="identityandview", executable_path=executable_path)
-```
+**NOTE:**: Installation outside a virtual environment may fail due to conflicting dependencies. Ensure compatibility, especially with `gym==0.21` and `numpy<=1.21.2`.
 
-6. For a full list of the NETT configs available, one can simply do the following. 
-```
-from netts.environment.configs import list_configs
-list_configs()
-```
+### Running a NETT
 
-Similar analogues for listing the encoders, policies, algorithms are also available: 
-```
-from netts.brain import list_algorithms, list_policies, list_encoders
-```
+1. **Download or Create the Unity Executable**: Obtain a pre-made Unity executable from [here](https://origins.luddy.indiana.edu/environments/). The executable is required to run the virtual environment.
 
-7. In order to orchestrate the benchmarking, all three components are brought together under one umbrella, the `NETT` object. This allows for storing details of the runs, serving as a reproducible artifact, automatic distribution of runs with different imprinting conditions, central logging, among other things. 
-```
-nett = NETT(brain=brain, body=body, environment=environment)
-```
+2. **Import NETT Components**:
+   Start by importing the NETT framework components - `Brain`, `Body`, and `Environment`, alongside the main `NETT` class.
+   ```python
+   from nett import Brain, Body, Environment
+   from nett import NETT
+   ```
 
-7. The created `nett` instance has the `.run()` method which carries out the execution. All the obvious runtime parameters with respect to benchmarking such as number of brains, train and test episodes, device(s) to be used for execution, steps per episode etc, are accepted as input to this function:
+3. **Component Configuration**:
+- **Brain**: Configure the learning aspects, including the policy network (e.g. "CnnPolicy"), learning algorithm (e.g. "PPO"), the reward function, and the encoder.
+   ```python
+   brain = Brain(policy="CnnPolicy", algorithm="PPO")
+   ```
+   To get a list of all available policies, algorithms, and encoders, the `Brain` class contains the methods `list_policies()`, `list_algorithms()`, and `list_encoders()` respectively.
+- **Body**: Set up the agent's physical interface with the environment. It's possible to apply gym.Wrappers for data preprocessing.
+   ```python
+   body = Body(type="basic", dvs=False, wrappers=None)
+   ```
+   Here, we do not pass any wrappers, letting information from the environment reach the brain "as is". Alternative body types (e.g. `two-eyed`, `rag-doll`) are planned in future updates.
 
-```
-run = netts.run(dir=run_dir, num_brains=5, trains_eps=10)
-```
+- **Environment**: Create the simulation environment using the path to your Unity executable (see Step 1).
+   ```python
+   environment = Environment(config="identityandview", executable_path="path/to/executable.x86_64")
+   ```
+   To get a list of all available configurations, run `Environment.list_configs()`.
 
-8. Internally, the function automatically creates a “task” list and assigns parallel processes for workers that execute the “jobs”. The result is the `run` object that can be used to inquire about the status of each of the “jobs” along with information about the device, mode etc. This is provided using the `.status()` method: 
-```
-nett.status(job_sheet)
-```
+4. **Run the Benchmarking**:  Integrate all components into a NETT instance to facilitate experiment execution.
+   ```python
+   benchmarks = NETT(brain=brain, body=body, environment=environment)
+   ```
+   The `NETT` instance has a `.run()` method that initiates the benchmarking process. The method accepts parameters such as the number of brains, training/testing episodes, and the output directory.
+   ```python
+   job_sheet = benchmarks.run(output_dir="path/to/run/output/directory/", num_brains=5, trains_eps=10, test_eps=5)
+   ```
+   The `run` function is asynchronous, returning the list of jobs that may or may not be complete.
+5. **Check Status**: To see the status of the benchmark processes, use the `.status()` method:
+   ```python
+   benchmarks.status(job_sheet)
+   ```
 
-Note that the `.run()` call executes in the background and does not block until its completion. This means the `run`` object is returned immediately and can be checked on a regular basis (programmatically or manually) to infer the completion of the run (or individual jobs).
+### Running Standard Analysis
 
+After running the experiments, the pipeline will generate a collection of datafiles in the defined output directory. 
 
-### **Running Standard Analysis**
 
-TO BE UPDATED.
+1. **Install R and dependencies**: To run the analyses performed in previous experiments,this toolkit provides a set of analysis scripts. Prior to running them, you will need R and the packages `tidyverse`, `argparse`, and `scales` installed. To install these packages, run the following command in R:
+   ```R
+   install.packages(c("tidyverse", "argparse", "scales"))
+   ```
+   Alternatively, if you are having difficulty installing R on your system, you can install these using conda.
+   ```bash
+   conda install -y r r-tidyverse r-argparse r-scales
+   ```
+2. **Run the Analysis**: To run the analysis, use the `analyze` method of the `NETT` class. This method will generate a set of plots and tables based on the datafiles in the output directory.
+   ```python
+   benchmarks.analyze(run_dir="path/to/run/output/directory/", output_dir="path/to/analysis/output/directory/")
+   ```
 
-## **Experiment Configuration**
+## Documentation
+For a link to the full documentation, please visit [here](https://buildingamind.github.io/NewbornEmbodiedTuringTest/).
+
+## Experiment Configuration
 
 More information related to details on the experiment can be found on following pages.
 
-* [**Parsing Experiment**](docs/Parsing.md)
-* [**ViewPoint Experiment**](docs/ViewInvariant.md)
\ No newline at end of file
+* [**Parsing Experiment**](docs/source/papers/Parsing.md)
+* [**ViewPoint Experiment**](docs/source/papers/ViewInvariant.md)
+
+[🔼 Back to top](#newborn-embodied-turing-test)
diff --git a/docs/Makefile b/docs/Makefile
new file mode 100644
index 00000000..d0c3cbf1
--- /dev/null
+++ b/docs/Makefile
@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
diff --git a/docs/assets/images/banner.png b/docs/assets/images/banner.png
new file mode 100644
index 00000000..edd7acec
Binary files /dev/null and b/docs/assets/images/banner.png differ
diff --git a/docs/digital_twin.jpg b/docs/assets/images/digital_twin.jpg
similarity index 100%
rename from docs/digital_twin.jpg
rename to docs/assets/images/digital_twin.jpg
diff --git a/docs/assets/images/digital_twin_cropped.jpg b/docs/assets/images/digital_twin_cropped.jpg
new file mode 100644
index 00000000..a5d1d1f0
Binary files /dev/null and b/docs/assets/images/digital_twin_cropped.jpg differ
diff --git a/docs/parsing.png b/docs/assets/images/parsing.png
similarity index 100%
rename from docs/parsing.png
rename to docs/assets/images/parsing.png
diff --git a/docs/viewpoint.png b/docs/assets/images/viewpoint.png
similarity index 100%
rename from docs/viewpoint.png
rename to docs/assets/images/viewpoint.png
diff --git a/docs/newbornembodied_0131_V1.docx b/docs/assets/newbornembodied_0131_V1.docx
similarity index 100%
rename from docs/newbornembodied_0131_V1.docx
rename to docs/assets/newbornembodied_0131_V1.docx
diff --git a/docs/parsing.mp4 b/docs/assets/videos/parsing.mp4
similarity index 100%
rename from docs/parsing.mp4
rename to docs/assets/videos/parsing.mp4
diff --git a/docs/viewpoint.mp4 b/docs/assets/videos/viewpoint.mp4
similarity index 100%
rename from docs/viewpoint.mp4
rename to docs/assets/videos/viewpoint.mp4
diff --git a/docs/make.bat b/docs/make.bat
new file mode 100644
index 00000000..747ffb7b
--- /dev/null
+++ b/docs/make.bat
@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.https://www.sphinx-doc.org/
+	exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd
diff --git a/docs/requirements.txt b/docs/requirements.txt
new file mode 100644
index 00000000..cfe9dffe
--- /dev/null
+++ b/docs/requirements.txt
@@ -0,0 +1,10 @@
+mlagents==1.0.0
+stable-baselines3[extra]==1.8.0
+sb3-contrib==1.8.0
+torchvision
+timm
+nvidia-ml-py
+sphinx
+mkdocs
+sphinx-rtd-theme
+myst-parser
\ No newline at end of file
diff --git a/docs/source/_static/video/parsing.mp4 b/docs/source/_static/video/parsing.mp4
new file mode 100644
index 00000000..6674e041
Binary files /dev/null and b/docs/source/_static/video/parsing.mp4 differ
diff --git a/docs/source/_static/video/viewpoint.mp4 b/docs/source/_static/video/viewpoint.mp4
new file mode 100644
index 00000000..7d838dce
Binary files /dev/null and b/docs/source/_static/video/viewpoint.mp4 differ
diff --git a/docs/source/conf.py b/docs/source/conf.py
new file mode 100644
index 00000000..8214a07a
--- /dev/null
+++ b/docs/source/conf.py
@@ -0,0 +1,44 @@
+'''Configuration file for the Sphinx documentation builder.'''
+
+# For the full list of built-in configuration values, see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+import os
+import sys
+sys.path.insert(0, os.path.abspath('../../src/'))
+
+source_suffix = {
+    '.rst': 'restructuredtext',
+    '.md': 'markdown',
+}
+
+# -- Project information -----------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
+
+project = 'NETT'
+copyright = '2024, Zachary Laborde'
+author = 'Zachary Laborde'
+release = '0.2'
+
+# -- General configuration ---------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
+
+extensions = [
+  'sphinx.ext.autodoc',
+  'myst_parser',
+]
+
+templates_path = ['_templates']
+exclude_patterns = []
+
+# -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
+
+html_theme = "sphinx_rtd_theme"
+html_theme_options = {
+    "body_max_width": "none",
+}
+
+html_static_path = ['_static']
+
+html_baseurl = "/html/"
\ No newline at end of file
diff --git a/docs/source/documentation/NETT.rst b/docs/source/documentation/NETT.rst
new file mode 100644
index 00000000..542ef4b5
--- /dev/null
+++ b/docs/source/documentation/NETT.rst
@@ -0,0 +1,5 @@
+NETT
+====
+
+.. autoclass:: nett.NETT
+  :members:
diff --git a/docs/source/documentation/body.rst b/docs/source/documentation/body.rst
new file mode 100644
index 00000000..1c7500ac
--- /dev/null
+++ b/docs/source/documentation/body.rst
@@ -0,0 +1,5 @@
+body
+====
+
+.. autoclass:: nett.Body
+   :members:
\ No newline at end of file
diff --git a/docs/source/documentation/brain.rst b/docs/source/documentation/brain.rst
new file mode 100644
index 00000000..792a7f60
--- /dev/null
+++ b/docs/source/documentation/brain.rst
@@ -0,0 +1,8 @@
+brain
+=====
+
+.. autoclass:: nett.Brain
+  :members:
+
+.. automodule:: nett.brain
+  :members:
diff --git a/docs/source/documentation/environment.rst b/docs/source/documentation/environment.rst
new file mode 100644
index 00000000..0c817571
--- /dev/null
+++ b/docs/source/documentation/environment.rst
@@ -0,0 +1,5 @@
+environment
+===========
+
+.. autoclass:: nett.Environment
+  :members:
diff --git a/docs/source/documentation/index.rst b/docs/source/documentation/index.rst
new file mode 100644
index 00000000..6ec31397
--- /dev/null
+++ b/docs/source/documentation/index.rst
@@ -0,0 +1,10 @@
+Documentation
+=============
+
+.. toctree:: 
+  :maxdepth: 4
+
+  brain
+  body
+  environment
+  NETT
diff --git a/docs/source/index.rst b/docs/source/index.rst
new file mode 100644
index 00000000..dfa115f4
--- /dev/null
+++ b/docs/source/index.rst
@@ -0,0 +1,40 @@
+.. NETT documentation master file, created by
+   sphinx-quickstart on Thu Feb  8 11:01:54 2024.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+===================================
+Newborn Embodied Turing Test (NETT)
+===================================
+
+Documentation
+=============
+
+.. toctree::
+   :titlesonly:
+
+   documentation/index
+   papers/index
+
+.. .. toctree::
+..    :titlesonly:
+
+..    documentation/NETT
+..    documentation/Brain
+..    documentation/Body
+..    documentation/Environment
+
+.. Research Publications
+.. =====================
+.. :ref:`papers/index`
+
+.. ..   papers/Parsing
+.. ..   papers/ViewInvariant
+
+Indices and tables
+==================
+
+* :ref:`genindex`
+
+.. * :ref:`search`
+.. * :ref:`modindex`
diff --git a/docs/Parsing.md b/docs/source/papers/Parsing.md
similarity index 84%
rename from docs/Parsing.md
rename to docs/source/papers/Parsing.md
index 6e3e89ab..5d98542d 100644
--- a/docs/Parsing.md
+++ b/docs/source/papers/Parsing.md
@@ -1,23 +1,24 @@
-## A newborn embodied Turing test for visual parsing
+# A Newborn Embodied Turing Test for Visual Parsing
 
 Manju Garimella, Denizhan Pak, Lalit Pandey, Justin N. Wood, & Samantha M. W. Wood
 
-[Parsing](https://github.com/buildingamind/pipeline_embodied/assets/1686251/839bd04c-8853-44c4-b275-5e61413a3904)
+<video controls>
+   <source src="../_static/video/parsing.mp4" type="video/mp4">
+</video>
 
-
-### Abstract
+## Abstract
 
 *Newborn brains exhibit remarkable abilities in rapid and generative learning, including the ability to parse objects from backgrounds and recognize those objects across substantial changes to their appearance (i.e., novel backgrounds and novel viewing angles). How can we build machines that can learn as efficiently as newborns? To accurately compare biological and artificial intelligence, researchers need to provide machines with the same training data that an organism has experienced since birth. Here, we present an experimental benchmark that enables researchers to raise artificial agents in the same controlled-rearing environments as newborn chicks. First, we raised newborn chicks in controlled environments with visual access to only a single object on a single background and tested their ability to recognize their object across novel viewing conditions. Then, we performed “digital twin” experiments in which we reared a variety of artificial neural networks in virtual environments that mimicked the rearing conditions of the chicks and measured whether they exhibited the same object recognition behavior as the newborn chicks. We found that biological chicks developed background-invariant object recognition, while the artificial chicks developed background-dependent recognition. Our benchmark exposes the limitations of current unsupervised and supervised algorithms in achieving the learning abilities of newborn animals. Ultimately, we anticipate that this approach will contribute to the development of AI systems that can learn with the same efficiency as newborn animals.*
 
-### Experiment Design
+## Experiment Design
 
 - VR chambers were equipped with two display walls (LCD monitors) for displaying object stimuli.
 - During the Training Phase, artificial chicks were reared in an environment containing a single 3D object rotating a full 360° around a horizontal axis in front of a naturalistic background scene. The object made a full rotation every 15s.
 - During the Test Phase, the VR chambers measured the artificial chicks’ imprinting response and object recognition performance. The “imprinting trials” measured whether the chicks developed an imprinting response.  The “test trials” measured the aritifical chicks’ ability to visually parse and recognize their imprinted object. During these trials, the imprinted object was presented on one display wall and an unfamiliar object was presented on the other display wall. Across the test trials, the objects were presented on all possible combinations of the three background scenes (Background 1 vs.Background 1, Background 1 vs. Background 2, Background 1 vs.Background 3, etc.).
 
-### Arguments
+## Arguments
 
-#### Train configuration
+### Train configuration
 
 ```
 agent_count: 1
@@ -42,27 +43,27 @@ Environment:
   record_agent: false
   recording_frames: 0
 ```
-##### Run script
+#### Run script
 
 ```bash
 python src/simulation/run_parsing_exp.py ++run_id=exp1 ++Environment.env_path=data/executables/parsing_benchmark/parsing_app.x86_64 ++mode=full ++train_eps=1000 ++test_eps=40 ++Agent.encoder="small" ++Environment.use_ship="true" ++Environment.background="A"
-
+```
 where
-Environment.use_ship = True or False (to choose between Ship and Fork);
-Environment.background = A, B, C (to choose between the three background);
+
+`Environment.use_ship` = True or False (to choose between Ship and Fork);
+`Environment.background` = A, B, C (to choose between the three background);
 mode = full or train or test (to choose between the three modes to run);
-Agent.encoder = "small", "medium" or "large" to choose between the three different types of encoders: NatureCNN, resnet10 and resnet18
-Agent.reward="supervised" default
+`Agent.encoder` = "small", "medium" or "large" to choose between the three different types of encoders: NatureCNN, resnet10 and resnet18
+`Agent.reward` = "supervised" default
 
-```
 
-##### Custom Configuration:
+#### Custom Configuration:
 
 - Update train episode count; test episode count
 - Encoder types - small, medium and large
 - Reward types - supervised or 'unsupervised'
 
-### Links
+## Links
 
-[Exectuables can be found here](https://origins.luddy.indiana.edu/unity/executables/)
+[Executables can be found here](https://origins.luddy.indiana.edu/unity/executables/)
 
diff --git a/docs/ViewInvariant.md b/docs/source/papers/ViewInvariant.md
similarity index 92%
rename from docs/ViewInvariant.md
rename to docs/source/papers/ViewInvariant.md
index e4bb63e8..66b24bb7 100644
--- a/docs/ViewInvariant.md
+++ b/docs/source/papers/ViewInvariant.md
@@ -1,25 +1,24 @@
-## A newborn embodied Turing test for View-Invariant Recognition
+# A Newborn Embodied Turing Test for View-Invariant Recognition
 
 Denizhan Pak, Donsuk Lee, Samantha M. W. Wood & Justin N. Wood
 
+<video controls>
+   <source src="../_static/video/viewpoint.mp4" type="video/mp4">
+</video>
 
-
-https://github.com/buildingamind/pipeline_embodied/assets/1686251/2fed4649-b4d6-4c93-813c-cd040a92c8cb
-
-
-### Abstract
+## Abstract
 
 *Recent progress in artificial intelligence has renewed interest in building machines that learn like animals. Almost all of the work comparing learning across biological and artificial systems comes from studies where animals and machines received different training data, obscuring whether differences between animals and machines emerged from differences in learning mechanisms versus training data. We present an experimental approach—a “newborn embodied Turing Test”—that allows newborn animals and machines to be raised in the same environments and tested with the same tasks, permitting direct comparison of their learning abilities. To make this platform, we first collected controlled-rearing data from newborn chicks, then performed “digital twin” experiments in which machines were raised in virtual environments that mimicked the rearing conditions of the chicks. We found that (1) machines (deep reinforcement learning agents with intrinsic motivation) can spontaneously develop visually guided preference behavior, akin to imprinting in newborn chicks, and (2) machines are still far from newborn-level performance on object recognition tasks. Almost all of the chicks developed view-invariant object recognition, whereas the machines tended to develop view-dependent recognition. The learning outcomes were also far more constrained in the chicks versus machines. Ultimately, we anticipate that this approach will help researchers develop embodied AI systems that learn like newborn animals.*
 
-### Experiment Design
+## Experiment Design
 
 - VR chambers were equipped with two display walls (LCD monitors) for displaying object stimuli.
 - During the Training Phase, artificial chicks were reared in an environment containing a single 3D object rotating 15° around a vertical axis in front of a blank background scene. The object made a full rotation every 3s. Agents can be imprinted to one of 4 possible conditions: side and front views of the Fork object or side and front views of the ship object.
 - During the Test Phase, the VR chambers measured the artificial chicks’ imprinting response and object recognition performance. The “imprinting trials” measured whether the chicks developed an imprinting response.  The “test trials” measured the aritifical chicks’ ability to visually discriminate their imprinted object. During these trials, the imprinted object, rotated at an alternate angle to the imprint condition, was presented on one display wall and an unfamiliar object was presented on the other display wall, the angle of which was either the same as the imprint condition (fixed trials) or matched to the viewpoint in the test condition (matched trials).
 
-### Arguments
+## Arguments
 
-#### Train configuration
+### Train configuration
 
 ```
 agent_count: 1
@@ -45,6 +44,6 @@ Environment:
   recording_frames: 0
 ```
 
-### Executables
+## Executables
 
 [Exectuable can be found here](https://origins.luddy.indiana.edu/unity/executables/).
diff --git a/docs/source/papers/index.rst b/docs/source/papers/index.rst
new file mode 100644
index 00000000..67f568b7
--- /dev/null
+++ b/docs/source/papers/index.rst
@@ -0,0 +1,7 @@
+Research Papers
+===============
+
+.. toctree::
+
+  Parsing
+  ViewInvariant
\ No newline at end of file
diff --git a/docs/index.md b/docs/source/papers/index0.md
similarity index 98%
rename from docs/index.md
rename to docs/source/papers/index0.md
index 82d7c851..ab6fa174 100644
--- a/docs/index.md
+++ b/docs/source/papers/index0.md
@@ -6,7 +6,7 @@ Lab**](http://buildingamind.com/)**. This pipeline provides all necessary compon
 
 The figure below shows the experiment setup for the three experiments discussed in the guide.
 
-<img src="digital_twin.jpg" alt="Digital Twin" style="zoom:35%;" />
+<img src="../../assets/images/digital_twin.jpg" alt="Digital Twin" style="zoom:35%;" />
 
 ## **How to Use this Repository**
 
diff --git a/mkdocs.yml b/mkdocs.yml
index bd7b9861..be833d44 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -1,4 +1,4 @@
-site_name: Embodied Pipeline Benchmarks
+site_name: Newborn Embodied Turing Test
 site_url: https://origin.luddy.indiana.edu
 nav:
     - Home: index.md
diff --git a/notebooks/Getting Started.ipynb b/notebooks/Getting Started.ipynb
index b9ae45c4..d7236249 100644
--- a/notebooks/Getting Started.ipynb
+++ b/notebooks/Getting Started.ipynb
@@ -7,8 +7,8 @@
    "outputs": [],
    "source": [
     "# imports\n",
-    "from netts import Brain, Body, Environment\n",
-    "from netts import NETT"
+    "from nett import Brain, Body, Environment\n",
+    "from nett import NETT"
    ]
   },
   {
@@ -42,7 +42,7 @@
     "# construct the NETT\n",
     "nett = NETT(brain=brain, body=body, environment=environment)\n",
     "# run the NETT\n",
-    "job_sheet = nett.run(dir=\"./test_run\", num_brains=1, train_eps=10, test_eps=1)"
+    "job_sheet = nett.run(output_dir=\"./test_run\", num_brains=1, train_eps=10, test_eps=1)"
    ]
   },
   {
diff --git a/pyproject.toml b/pyproject.toml
index 904850ca..8d8fe50d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,25 +4,41 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "nett-benchmarks"
-version = "0.1"
+dynamic = ["version"]
 authors = [
   { name="Bhargav Desai", email="desabh@iu.edu" },
+  { name="Zachary Laborde", email="zlaborde@iu.edu" },
 ]
 description = "A testbed for comparing the learning abilities of newborn animals and autonomous artificial agents."
 readme = "README.md"
-requires-python = ">=3.10.12, <3.10.13"
+license = {file = "LICENSE"}
+requires-python = "==3.10.12"
 dependencies = [
     "mlagents==1.0.0",
     "stable-baselines3[extra]==1.8.0",
     "sb3-contrib==1.8.0",
-    "imageio-ffmpeg",
     "torchvision",
-    "scipy",
     "timm",
     "nvidia-ml-py",
-    "ipywidgets"
 ]
+keywords = ["nett","netts","newborn","embodied","turing test","benchmark","benchmarking","learning","animals","autonomous","artificial","agents","reinforcement","neuroml","AI","ML","machine learning","artificial intelligence"]
+classifiers = [
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: R",
+  "Environment :: GPU :: NVIDIA CUDA"
+]
+
+[project.optional-dependencies]
+notebook = ["ipywidgets"]
 
 [project.urls]
-"Homepage" = "https://github.com/pypa/sampleproject"
-"Bug Tracker" = "https://github.com/pypa/sampleproject/issues"
\ No newline at end of file
+"Homepage" = "https://buildingamind.github.io/NewbornEmbodiedTuringTest/"
+"Bug Tracker" = "https://buildingamind.github.io/NewbornEmbodiedTuringTest/issues"
+"Documentation" = "https://buildingamind.github.io/NewbornEmbodiedTuringTest/index.html"
+
+[tool.setuptools.dynamic]
+version = { attr = "nett.__version__" }
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
diff --git a/src/netts/__init__.py b/src/nett/__init__.py
similarity index 56%
rename from src/netts/__init__.py
rename to src/nett/__init__.py
index fe20ca20..22387ff7 100644
--- a/src/netts/__init__.py
+++ b/src/nett/__init__.py
@@ -12,20 +12,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+'''
+Initialize the NETT library
+'''
+
+import os
 import logging
 from pathlib import Path
 # simplify imports
-from netts.brain.builder import Brain
-from netts.body.builder import Body
-from netts.environment.builder import Environment
-from netts.nett import NETT
+from nett.brain.builder import Brain
+from nett.body.builder import Body
+from nett.environment.builder import Environment
+from nett.nett import NETT
 
 # release version
-__version__ = "0.1"
+__version__ = "0.2"
+
+# change permissions of the ml-agents binaries directory
+os.chmod('/tmp/ml-agents-binaries', 0o1777)
+os.chmod('/tmp/ml-agents-binaries/binaries', 0o1777)
+os.chmod('/tmp/ml-agents-binaries/tmp', 0o1777)
 
 # path to store library cache (such as configs etc)
-cache_dir = Path.joinpath(Path.home(), '.cache', 'netts')
+cache_dir = Path.joinpath(Path.home(), ".cache", "nett")
 
 # set up logging
-logging.basicConfig(format='[%(name)s] %(levelname)s:  %(message)s', level=logging.INFO)
-logger = logging.getLogger(__name__)
\ No newline at end of file
+logging.basicConfig(format="[%(name)s] %(levelname)s:  %(message)s", level=logging.INFO)
+logger = logging.getLogger(__name__)
diff --git a/src/netts/utils/__init__.py b/src/nett/analysis/.Rapp.history
similarity index 100%
rename from src/netts/utils/__init__.py
rename to src/nett/analysis/.Rapp.history
diff --git a/src/nett/analysis/.Rhistory b/src/nett/analysis/.Rhistory
new file mode 100644
index 00000000..cc92306d
--- /dev/null
+++ b/src/nett/analysis/.Rhistory
@@ -0,0 +1,482 @@
+favorite_numbers <- c(50, 3, 7, 2, 8, 4, 4, 7, 13, 7, 22, 23, 7, 30, 4, 7, 12)
+table(favorite_numbers)
+mode(favorite_numbers)
+summary(favorite_numbers)
+mode.sngl(favorite_numbers)
+mode.mult(favorite_numbers)
+favorite_foods <- c(“pizza”, “nachos”, “salad”, “pizza”, “pizza”, “tacos”, “guacamole”, “pasta”, “soup”, “salad”, “nachos”, “rice”, “pizza”, “tacos”)
+favorite_foods <- c(“pizza”, “nachos”, “salad”, “pizza”, “pizza”, “tacos”, “guacamole”, “pasta”, “soup”, “salad”, “nachos”, “rice”, “pizza”, “tacos”)
+bananas <- 225
+bananas * 365
+bananas
+5+2
+9-6
+2*4
+6/2
+4^2
+346+975
+37*56
+1440/24
+17^4
+5+2
+# Subtraction
+9-6
+range(3,4,5,6,9)
+favorite_foods <- c(“pizza”, “nachos”, “salad”, “pizza”, “pizza”, “tacos”, “guacamole”, “pasta”, “soup”, “salad”, “nachos”, “rice”, “pizza”, “tacos”)
+favorite_foods <- c(“pizza”, “nachos”, “salad”, “pizza”, “pizza”, “tacos”, “guacamole”, “pasta”, “soup”, “salad”, “nachos”, “rice”, “pizza”, “tacos”)
+favorite_foods <- c("pizza", "nachos", "salad")
+summary(favorite_foods)
+numbers <- c(1,2,3)
+summary(numbers)
+5+2
+9-6
+2*4
+6/2
+4^2
+5+2
+9-6
+2*4
+6/2
+4^2
+5*82/13
+mean(2,3,4,6,834,9)
+paper <- 135
+ink <- 15
+rev_p <- 5
+rev_i <- 10
+(paper*rev_p) + (ink*rev_i)
+march_sales <- c(835, 680, 930, 870, 815, 890, 700, 925, 830, 890)
+mean(march_sales)
+median(march_sales)
+table(march_sales)
+my_tab <- table(march_sales)
+my_tab
+max(my_tab)
+my_tab[max(my_tab)]
+march_sales <- c(835, 680, 930, 870, 815, 890, 700, 925, 830, 890)
+hist(march_sales)
+hist(march_sales, xlab="")
+hist(march_sales, xlab="", breaks=6)
+hist(march_sales, xlab="", breaks=5)
+hist(march_sales, xlab="", breaks = 10
+)
+hist(march_sales, xlab="", breaks = c(600,650,700,750,800,850,900,950)
+)
+5*8/13
+paper <- 125
+paper * 365
+paper - 10
+paper <- 125
+paper * 365
+paper - 10
+# Central Tendency
+mean(8,6,7,5,3,0,9)
+median(8,6,7,5,3,0,9)
+# Basic Math with Variables
+paper <- 125
+mode(8,6,7,5,3,0,9)
+table(8,6,7,5,3,0,9)
+table(8,6,7,5,3,0,9)
+vals <- c(8,6,7,5,3,0,9)
+mean(vals)
+jenny <- c(8,6,7,5,3,0,9)
+# Central Tendency
+mean(jenny)
+median(jenny)
+table(jenny)
+march_sales <- c(835, 680, 930, 870, 815, 890, 700, 925, 830, 890)
+# Measures of spread
+max(march_sales) -min(march_sales)
+sd(march_sales)
+IQR(march_sales)
+MARCH_SALES
+summary(march_sales)
+# Histogram
+hist(march_sales)
+hist(march_sales, main="Dunder Mifflin March Sales")
+hist(march_sales,
+main="Dunder Mifflin March Sales",
+xlab="")
+hist(march_sales,
+main="Dunder Mifflin March Sales",
+xlab="",
+breaks=3)
+hist(march_sales,
+main="Dunder Mifflin March Sales",
+xlab="",
+breaks=20)
+hist(march_sales,
+main="Dunder Mifflin March Sales",
+xlab="",
+breaks=4)
+hist(march_sales,
+main="Dunder Mifflin March Sales",
+xlab="",
+breaks=4,
+ylim=c(0,4))
+hist(march_sales,
+main="Dunder Mifflin March Sales",
+xlab="",
+breaks=4,
+ylim=c(0,4),
+col="red")
+hist(march_sales,
+main="Dunder Mifflin March Sales",
+xlab="",
+breaks=4,
+ylim=c(0,4),
+col="turquoise")
+colors()
+# Histogram
+hist(march_sales)
+hist(march_sales,
+main="Dunder Mifflin March Sales",
+xlab="",
+breaks=4,
+ylim=c(0,4),
+col="turquoise")
+lkdlfndf
+ljbdfjsdf
+data<-c(8,6,7,5,3,0,9)
+boxplot(data)
+boxplot(data, main="My awesome title!", col="turquoise")
+evermore <- c(533, 296, 161, 133, 129, 158, 106, 92, 100, 116, 94, 108, 87, 69, 114, 121, 52)
+mean(evermore)
+median(evermore)
+summary(evermore)
+boxplot(evermore)
+midnights <- c(86, 75, 128, 80, 71, 78, 62, 65, 65, 51, 67, 49, 51, 67, 49, 51, 29, 29, 26, 24, 22, 32, 21)
+evermore <- c(533, 296, 161, 133, 129, 158, 106, 92, 100, 116, 94, 108, 87, 69, 114, 121, 52)
+mean(evermore)
+median(evermore)
+library(pwr)
+install.packages(pwr)
+install.packages("pwr")
+library(pwr)
+# These are all 12-fold results (not sparse)
+simclr_fast <- c(.870, .830, .810, .810, .800, .855)
+simclr_slow <- c(.850, .852, .870, .886, .842, .890)
+simclr_fast_ns <- c(.850, .855, .830, .820, .810, .795)
+simclr_slow_ns <- c(.782, .865, .835, .840, .820, .875)
+vit_fast <- c(.736, .750, .720, .726, .720, .730)
+vit_slow <- c(.780, .762, .783, .775, .761, .735)
+meandiff_simclr_slow_vs_fast <- abs( mean(simclr_fast) - mean(simclr_slow) )
+sd_simclr_slow_vs_fast <- sqrt( (sd(simclr_fast)^2 + sd(simclr_slow)^2) /2 )
+d_simclr_slow_vs_fast <- meandiff_simclr_slow_vs_fast / sd_simclr_slow_vs_fast
+sd(simclr_fast)
+sd(simclr_slow)
+sd_simclr_slow_vs_fast <- sqrt( (2*sd(simclr_fast)^2 + 2*sd(simclr_slow)^2) / 4)
+sd_simclr_slow_vs_fast <- sqrt( (2*sd(simclr_fast)^2 + 2*sd(simclr_slow)^2) / 4)
+d_simclr_slow_vs_fast <- meandiff_simclr_slow_vs_fast / sd_simclr_slow_vs_fast
+meandiff_simclr_slow_vs_ns_slow <- abs( mean(simclr_slow_ns) - mean(simclr_slow) )
+sd_simclr_slow_vs_ns_slow <- sqrt( (2*sd(simclr_slow_ns)^2 + 2*sd(simclr_slow)^2) / 4)
+d_simclr_slow_vs_ns_slow <- meandiff_simclr_slow_vs_ns_slow / sd_simclr_slow_vs_ns_slow
+meandiff_simclr_slow_vs_ns_fast <- abs( mean(simclr_fast_ns) - mean(simclr_slow) )
+sd_simclr_slow_vs_ns_fast <- sqrt( (2*sd(simclr_fast_ns)^2 + 2*sd(simclr_slow)^2) / 4)
+d_simclr_slow_vs_ns_fast <- meandiff_simclr_slow_vs_ns_fast / sd_simclr_slow_vs_ns_fast
+### vit - slow vs. fast
+meandiff_vit_slow_vs_fast <- abs( mean(vit_fast) - mean(vit_slow) )
+sd_vit_slow_vs_fast <- sqrt( (2*sd(vit_fast)^2 + 2*sd(vit_slow)^2) / 4)
+d_vit_slow_vs_fast <- meandiff_vit_slow_vs_fast / sd_vit_slow_vs_fast
+pwr_simclr_slow_vs_fast <- pwr.t.test(d=d_simclr_slow_vs_fast, sig.level=.05, power=power,
+type="two.sample", alternative="two.sided")
+power = 0.8
+pwr_simclr_slow_vs_fast <- pwr.t.test(d=d_simclr_slow_vs_fast, sig.level=.05, power=power,
+type="two.sample", alternative="two.sided")
+View(pwr_simclr_slow_vs_fast)
+pwr_simclr_slow_vs_ns_slow <- pwr.t.test(d=d_simclr_slow_vs_ns_slow, sig.level=.05, power=power,
+type="two.sample", alternative="two.sided")
+pwr_simclr_slow_vs_ns_fast <- pwr.t.test(d=d_simclr_slow_vs_ns_fast, sig.level=.05, power=power,
+type="two.sample", alternative="two.sided")
+pwr_vit_slow_vs_fast <- pwr.t.test(d=d_vit_slow_vs_fast, sig.level=.05, power=power,
+type="two.sample", alternative="two.sided")
+results <- df(condition=c("simclr slow vs fast", "simclr slow vs non-smooth slow",
+"simclr slow vs non-smooth fast", "vit slow vs fast"),
+mean_diff=c(meandiff_simclr_slow_vs_fast, meandiff_simclr_slow_vs_ns_slow,
+meandiff_simclr_slow_vs_ns_fast, meandiff_vit_slow_vs_fast),
+pooled_sd=c(sd_simclr_slow_vs_fast, sd_simclr_slow_vs_ns_slow,
+sd_simclr_slow_vs_ns_fast, sd_vit_slow_vs_fast),
+cohens_d=c(d_simclr_slow_vs_fast, d_simclr_slow_vs_ns_slow,
+sd_simclr_slow_vs_ns_fast, d_vit_slow_vs_fast),
+sample_for_power = c(pwr_simclr_slow_vs_fast$n, pwr_simclr_slow_vs_ns_slow$n,
+pwr_simclr_slow_vs_ns_fast$n, pwr_vit_slow_vs_fast$n))
+results <- data.frame(condition=c("simclr slow vs fast", "simclr slow vs non-smooth slow",
+"simclr slow vs non-smooth fast", "vit slow vs fast"),
+mean_diff=c(meandiff_simclr_slow_vs_fast, meandiff_simclr_slow_vs_ns_slow,
+meandiff_simclr_slow_vs_ns_fast, meandiff_vit_slow_vs_fast),
+pooled_sd=c(sd_simclr_slow_vs_fast, sd_simclr_slow_vs_ns_slow,
+sd_simclr_slow_vs_ns_fast, sd_vit_slow_vs_fast),
+cohens_d=c(d_simclr_slow_vs_fast, d_simclr_slow_vs_ns_slow,
+sd_simclr_slow_vs_ns_fast, d_vit_slow_vs_fast),
+sample_for_power = c(pwr_simclr_slow_vs_fast$n, pwr_simclr_slow_vs_ns_slow$n,
+pwr_simclr_slow_vs_ns_fast$n, pwr_vit_slow_vs_fast$n))
+results
+# Slow & Smooth Power Analysis
+# Comparisons:
+### slow vs fast
+### slow vs non-smooth slow
+### slow vs non-smooth fast
+install.packages("pwr")
+library(pwr)
+power = 0.8
+##################################################################
+########################       DATA       ########################
+##################################################################
+# These are all 12-fold results (not sparse)
+simclr_fast <- c(.870, .830, .810, .810, .800, .855)
+simclr_slow <- c(.850, .852, .870, .886, .842, .890)
+simclr_fast_ns <- c(.850, .855, .830, .820, .810, .795)
+simclr_slow_ns <- c(.782, .865, .835, .840, .820, .875)
+vit_fast <- c(.736, .750, .720, .726, .720, .730)
+vit_slow <- c(.780, .762, .783, .775, .761, .735)
+##################################################################
+######################       Cohen's D       #####################
+##################################################################
+# Compute Effect Size
+### simclr - slow vs. fast
+meandiff_simclr_slow_vs_fast <- abs( mean(simclr_fast) - mean(simclr_slow) )
+sd_simclr_slow_vs_fast <- sqrt( (2*sd(simclr_fast)^2 + 2*sd(simclr_slow)^2) / 4)
+d_simclr_slow_vs_fast <- meandiff_simclr_slow_vs_fast / sd_simclr_slow_vs_fast
+### simclr - slow vs. non-smooth slow
+meandiff_simclr_slow_vs_ns_slow <- abs( mean(simclr_slow_ns) - mean(simclr_slow) )
+sd_simclr_slow_vs_ns_slow <- sqrt( (2*sd(simclr_slow_ns)^2 + 2*sd(simclr_slow)^2) / 4)
+d_simclr_slow_vs_ns_slow <- meandiff_simclr_slow_vs_ns_slow / sd_simclr_slow_vs_ns_slow
+### simclr - slow vs. non-smooth fast
+meandiff_simclr_slow_vs_ns_fast <- abs( mean(simclr_fast_ns) - mean(simclr_slow) )
+sd_simclr_slow_vs_ns_fast <- sqrt( (2*sd(simclr_fast_ns)^2 + 2*sd(simclr_slow)^2) / 4)
+d_simclr_slow_vs_ns_fast <- meandiff_simclr_slow_vs_ns_fast / sd_simclr_slow_vs_ns_fast
+### vit - slow vs. fast
+meandiff_vit_slow_vs_fast <- abs( mean(vit_fast) - mean(vit_slow) )
+sd_vit_slow_vs_fast <- sqrt( (2*sd(vit_fast)^2 + 2*sd(vit_slow)^2) / 4)
+d_vit_slow_vs_fast <- meandiff_vit_slow_vs_fast / sd_vit_slow_vs_fast
+##################################################################
+########################       POWER       #######################
+##################################################################
+pwr_simclr_slow_vs_fast <- pwr.t.test(d=d_simclr_slow_vs_fast, sig.level=.05, power=power,
+type="two.sample", alternative="two.sided")
+pwr_simclr_slow_vs_ns_slow <- pwr.t.test(d=d_simclr_slow_vs_ns_slow, sig.level=.05, power=power,
+type="two.sample", alternative="two.sided")
+pwr_simclr_slow_vs_ns_fast <- pwr.t.test(d=d_simclr_slow_vs_ns_fast, sig.level=.05, power=power,
+type="two.sample", alternative="two.sided")
+pwr_vit_slow_vs_fast <- pwr.t.test(d=d_vit_slow_vs_fast, sig.level=.05, power=power,
+type="two.sample", alternative="two.sided")
+##################################################################
+########################       RESULTS       #####################
+##################################################################
+results <- data.frame(condition=c("simclr slow vs fast", "simclr slow vs non-smooth slow",
+"simclr slow vs non-smooth fast", "vit slow vs fast"),
+mean_diff=c(meandiff_simclr_slow_vs_fast, meandiff_simclr_slow_vs_ns_slow,
+meandiff_simclr_slow_vs_ns_fast, meandiff_vit_slow_vs_fast),
+pooled_sd=c(sd_simclr_slow_vs_fast, sd_simclr_slow_vs_ns_slow,
+sd_simclr_slow_vs_ns_fast, sd_vit_slow_vs_fast),
+cohens_d=c(d_simclr_slow_vs_fast, d_simclr_slow_vs_ns_slow,
+d_simclr_slow_vs_ns_fast, d_vit_slow_vs_fast),
+sample_for_power = c(pwr_simclr_slow_vs_fast$n, pwr_simclr_slow_vs_ns_slow$n,
+pwr_simclr_slow_vs_ns_fast$n, pwr_vit_slow_vs_fast$n))
+print(results)
+install.packages("argparse")
+library(ggsci)
+install.packages("ggsci")
+library(ggsci)
+install.packages("RColorBrewer")
+paper <- 125
+class_numbers <- c(3, 12, 4, 13, 5, 19, 17, 13)
+class_numbers
+class_pets <- c("dog", "dog
+class_pets <- c("dog", "dog")
+mean(class_numbers)
+median(class_numbers)
+mode(class_numbers)
+table(class_numbers)
+min(class_numbers)
+max(class_numbers)
+sd(class_numbers)
+IQR(class_numbers)
+summary(class_numbers)
+summary(class_pets)
+hist(class_numbers)
+hist(class_numbers, "Class Numbers")
+hist(class_numbers, main="Class Numbers")
+hist(class_numbers, main="Class Numbers", xlab="Class Favorite Numbers")
+hist(class_numbers,
+main="Class Numbers",
+xlab="Class Favorite Numbers")
+hist(class_numbers,
+main="Class Numbers",
+xlab="Class Favorite Numbers",
+breaks=c(0,3,6,9,12,15,18,21))
+class_numbers <- c(3, 12, 4, 13, 5, 19, 17, 13, 3, 6, 2)
+hist(class_numbers,
+main="Class Numbers",
+xlab="Class Favorite Numbers",
+breaks=c(0,3,6,9,12,15,18,21))
+hist(class_numbers,
+main="Class Numbers",
+xlab="Class Favorite Numbers",
+breaks=c (0,3,6,9,12,15,18,21))
+hist(class_numbers,
+main = "Class Numbers",
+xlab = "Class Favorite Numbers",
+breaks = c(0, 3, 6, 9, 12, 15, 18, 21))
+hist(class_numbers,
+main = "Class Numbers",
+xlab = "Class Favorite Numbers",
+breaks = c(0, 3, 6, 9, 12, 15, 18, 21),
+col = "blue")
+source("NETT_test_viz.R")
+setwd("/Users/samanthawood/Documents/WoodLab/AI_Papers/EmbodiedPipeline/EmbodiedAnalysis")
+source("NETT_test_viz.R")
+graphstyle <- "/Users/samanthawood/Documents/WoodLab/AI_Papers/EmbodiedPipeline/EmbodiedAnalysis/NETT_graphstyle.R"
+View(parser)
+graphstyle <- "/Users/samanthawood/Documents/WoodLab/AI_Papers/EmbodiedPipeline/EmbodiedAnalysis/NETT_graphstyle.R"
+source(graphstyle)
+library(tidyverse)
+source(graphstyle)
+chick_data <- '/Users/samanthawood/Documents/WoodLab/AI_Papers/EmbodiedPipeline/EmbodiedAnalysis/ChickData/ChickData_Parsing.csv'
+# Load the chick data
+chick_data <- read.csv(chick_file)
+# Load the chick data
+chick_data <- read.csv(chick_data)
+View(chick_data)
+chick_data$cond_name <- gsub("LB", "\n", chick_data$cond_name)
+View(chick_data)
+chick_file <- '/Users/samanthawood/Documents/WoodLab/AI_Papers/EmbodiedPipeline/EmbodiedAnalysis/ChickData/ChickData_Parsing.csv'
+# Load the chick data
+chick_data <- read.csv(chick_file)
+chick_data$cond_name <- gsub("LB", "\n", chick_data$cond_name)
+View(chick_data)
+# Plot aesthetic settings ------------------------------------------------------
+custom_palette <- c("#3F8CB7", "#FCEF88", "#5D5797", "#62AC6B", "#B74779")
+chickred <- "#AF264A"
+p <- ggplot() +
+theme_classic() +
+theme(axis.text.x = element_text(size = 6)) +
+ylab("Percent Correct") +
+xlab("Test Condition") +
+scale_y_continuous(expand = c(0, 0), limits = c(0, 1), breaks=seq(0,1,.1), labels = scales::percent) +
+geom_hline(yintercept = .5, linetype = 2) +
+scale_fill_manual(values = custom_palette) +
+scale_colour_manual(values = custom_palette) +
+theme(axis.title = element_text(face="bold"),
+axis.text.x = element_text(face="bold", size=8.5),
+axis.text.y = element_text(face="bold", size=8.5))
+rm(make_bar_charts())
+rm(make_bar_charts
+)
+color_bars <- TRUE
+View(chick_data)
+p +
+# Model performance: bars
+{if(color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg, fill = cond_name))}+
+{if(!color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg), fill = "gray45")}+
+# Model performance: error bars
+geom_errorbar(data = chick_data, width = 0.3,
+aes(x = cond_name, ymin = avg-se, ymax = avg+se)) +
+ggsave(img_name, width = 6, height = 6)
+img_name <- "chick_results_parsing.png"
+p +
+# Model performance: bars
+{if(color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg, fill = cond_name))}+
+{if(!color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg), fill = "gray45")}+
+# Model performance: error bars
+geom_errorbar(data = chick_data, width = 0.3,
+aes(x = cond_name, ymin = avg-se, ymax = avg+se))
+ggsave(img_name, width = 6, height = 6)
+p +
+# Model performance: bars
+{if(color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg, fill = cond_name))}+
+{if(!color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg), fill = "gray45")}+
+# Model performance: error bars
+geom_errorbar(data = chick_data, width = 0.3,
+aes(x = cond_name, ymin = avg-se, ymax = avg+se)) +
+theme(legend.position="none")
+ggsave(img_name, width = 6, height = 6)
+original_order <- unique(chick_data$cond_name)
+chick_data$cond_name <- factor(chick_data$cond_name, levels = original_order)
+p +
+# Model performance: bars
+{if(color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg, fill = cond_name))}+
+{if(!color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg), fill = "gray45")}+
+# Model performance: error bars
+geom_errorbar(data = chick_data, width = 0.3,
+aes(x = cond_name, ymin = avg-se, ymax = avg+se)) +
+theme(legend.position="none")
+ggsave(img_name, width = 6, height = 6)
+# Plot aesthetic settings ------------------------------------------------------
+custom_palette <- c("#FCEF88", "#5D5797", "#3F8CB7", "#62AC6B", "#B74779")
+p <- ggplot() +
+theme_classic() +
+theme(axis.text.x = element_text(size = 6)) +
+ylab("Percent Correct") +
+xlab("Test Condition") +
+scale_y_continuous(expand = c(0, 0), limits = c(0, 1), breaks=seq(0,1,.1), labels = scales::percent) +
+geom_hline(yintercept = .5, linetype = 2) +
+scale_fill_manual(values = custom_palette) +
+scale_colour_manual(values = custom_palette) +
+theme(axis.title = element_text(face="bold"),
+axis.text.x = element_text(face="bold", size=8.5),
+axis.text.y = element_text(face="bold", size=8.5))
+# Make that plot! --------------------------------------------------------------
+p +
+# Model performance: bars
+{if(color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg, fill = cond_name))}+
+{if(!color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg), fill = "gray45")}+
+# Model performance: error bars
+geom_errorbar(data = chick_data, width = 0.3,
+aes(x = cond_name, ymin = avg-se, ymax = avg+se)) +
+theme(legend.position="none")
+ggsave(img_name, width = 6, height = 6)
+chick_by_subj <- /Users/samanthawood/Documents/WoodLab/AI_Papers/EmbodiedPipeline/EmbodiedAnalysis/ChickData/chicks_by_subj.csv
+chick_by_subj <- "/Users/samanthawood/Documents/WoodLab/AI_Papers/EmbodiedPipeline/EmbodiedAnalysis/ChickData/chicks_by_subj.csv"
+# Load the chick data (by subject)
+by_subj <- read.csv(chick_by_subj)
+View(by_subj)
+# Make that plot! --------------------------------------------------------------
+p +
+# Model performance: bars
+{if(color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg, fill = cond_name))}+
+{if(!color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg), fill = "gray45")}+
+# Model performance: error bars
+geom_errorbar(data = chick_data, width = 0.3,
+aes(x = cond_name, ymin = avg-se, ymax = avg+se)) +
+geom_jitter(data=by_subj, aes(x=cond_name, y=Performance)) +
+theme(legend.position="none")
+ggsave(img_name, width = 6, height = 6)
+by_subj$cond_name <- gsub("LB", "\n", by_subj$cond_name)
+View(by_subj)
+# Plot aesthetic settings ------------------------------------------------------
+custom_palette <- c("#FCEF88", "#5D5797", "#3F8CB7", "#62AC6B", "#B74779")
+chickred <- "#AF264A"
+p <- ggplot() +
+theme_classic() +
+theme(axis.text.x = element_text(size = 6)) +
+ylab("Percent Correct") +
+xlab("Test Condition") +
+scale_y_continuous(expand = c(0, 0), limits = c(0, 1), breaks=seq(0,1,.1), labels = scales::percent) +
+geom_hline(yintercept = .5, linetype = 2) +
+scale_fill_manual(values = custom_palette) +
+scale_colour_manual(values = custom_palette) +
+theme(axis.title = element_text(face="bold"),
+axis.text.x = element_text(face="bold", size=8.5),
+axis.text.y = element_text(face="bold", size=8.5))
+# Make that plot! --------------------------------------------------------------
+p +
+# Model performance: bars
+{if(color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg, fill = cond_name))}+
+{if(!color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg), fill = "gray45")}+
+# Model performance: error bars
+geom_errorbar(data = chick_data, width = 0.3,
+aes(x = cond_name, ymin = avg-se, ymax = avg+se)) +
+geom_jitter(data=by_subj, aes(x=cond_name, y=Performance)) +
+theme(legend.position="none")
+ggsave(img_name, width = 6, height = 6)
+p +
+# Model performance: bars
+{if(color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg, fill = cond_name))}+
+{if(!color_bars)geom_col(data = chick_data, width = 0.7, aes(x=cond_name, y = avg), fill = "gray45")}+
+# Model performance: error bars
+geom_errorbar(data = chick_data, width = 0.3,
+aes(x = cond_name, ymin = avg-se, ymax = avg+se)) +
+geom_jitter(data=by_subj, aes(x=cond_name, y=Performance), width = .3) +
+theme(legend.position="none")
+ggsave(img_name, width = 6, height = 6)
diff --git a/src/nett/analysis/.gitattributes b/src/nett/analysis/.gitattributes
new file mode 100644
index 00000000..dfe07704
--- /dev/null
+++ b/src/nett/analysis/.gitattributes
@@ -0,0 +1,2 @@
+# Auto detect text files and perform LF normalization
+* text=auto
diff --git a/src/nett/analysis/ChickData/.DS_Store b/src/nett/analysis/ChickData/.DS_Store
new file mode 100644
index 00000000..5008ddfc
Binary files /dev/null and b/src/nett/analysis/ChickData/.DS_Store differ
diff --git a/src/nett/analysis/ChickData/ChickData_Parsing.csv b/src/nett/analysis/ChickData/ChickData_Parsing.csv
new file mode 100644
index 00000000..46f1ae4a
--- /dev/null
+++ b/src/nett/analysis/ChickData/ChickData_Parsing.csv
@@ -0,0 +1,5 @@
+test.cond,avg,std_dev,avg_dev,se
+Imprinted Object Familiar,0.652731096,0.118585046,0.091188669,0.021298503
+Novel Familiar,0.722057353,0.100726968,0.081165209,0.018091098
+Both Familiar,0.707892313,0.109402977,0.087105549,0.019649355
+Both Unfamiliar,0.725346606,0.101022594,0.084554842,0.018144193
diff --git a/src/nett/analysis/ChickData/chicks_by_subj.csv b/src/nett/analysis/ChickData/chicks_by_subj.csv
new file mode 100644
index 00000000..9133da0b
--- /dev/null
+++ b/src/nett/analysis/ChickData/chicks_by_subj.csv
@@ -0,0 +1,125 @@
+Subject,Performance,test.cond
+o1_highway_C1,0.823888116,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_highway_C2,0.774409362,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_highway_C3,0.878543678,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_highway_C4,0.666174972,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_highway_C5,0.728370471,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_highway_C1,0.893385664,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_highway_C2,0.64886525,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_highway_C3,0.628141675,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_highway_C4,0.659863177,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_highway_C5,0.726226164,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_highway_C6,0.685786015,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_beach_C1,0.683451016,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_beach_C2,0.490077173,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_beach_C3,0.646481964,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_beach_C4,0.883947655,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_beach_C5,0.797743266,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_beach_C6,0.883002114,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_beach_C1,0.650247204,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_beach_C2,0.749652739,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_beach_C3,0.821864978,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_beach_C4,0.802221107,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_beach_C5,0.651864285,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_beach_C6,0.800491686,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_forest_C1,0.797924512,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_forest_C2,0.806411838,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_forest_C3,0.74806239,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_forest_C4,0.597519634,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_forest_C5,0.605336908,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_forest_C3,0.657043339,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_forest_C4,0.709763154,Familiar Object Novel Background vs. Novel Object Novel Background
+o2_forest_C6,0.588983276,Familiar Object Novel Background vs. Novel Object Novel Background
+o1_highway_C1,0.85506525,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_highway_C2,0.654969338,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_highway_C3,0.849836637,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_highway_C4,0.771973881,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_highway_C5,0.631057839,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_highway_C1,0.906921221,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_highway_C2,0.666661105,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_highway_C3,0.582380658,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_highway_C4,0.579476732,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_highway_C5,0.668244069,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_highway_C6,0.683952788,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_beach_C1,0.696865909,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_beach_C2,0.595385804,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_beach_C3,0.500444609,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_beach_C4,0.908621707,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_beach_C5,0.711471374,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_beach_C6,0.938951245,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_beach_C1,0.697306775,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_beach_C2,0.724868742,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_beach_C3,0.802479279,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_beach_C4,0.723031214,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_beach_C5,0.644041126,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_beach_C6,0.825637946,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_forest_C1,0.729598913,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_forest_C2,0.736411428,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_forest_C3,0.775759549,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_forest_C4,0.578463958,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_forest_C5,0.616029562,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_forest_C3,0.660968059,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_forest_C4,0.633527107,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o2_forest_C6,0.594257884,Familiar Object Familiar Background vs. Novel Object Familiar Background
+o1_highway_C1,0.684461007,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_highway_C2,0.689396456,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_highway_C3,0.833442132,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_highway_C4,0.578204834,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_highway_C5,0.705602922,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_highway_C1,0.82813645,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_highway_C2,0.688919158,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_highway_C3,0.349914442,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_highway_C4,0.652912353,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_highway_C5,0.581763919,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_highway_C6,0.556956447,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_beach_C1,0.635776349,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_beach_C2,0.651773253,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_beach_C3,0.45581979,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_beach_C4,0.836997262,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_beach_C5,0.672052859,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_beach_C6,0.79843297,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_beach_C1,0.56308097,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_beach_C2,0.719796476,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_beach_C3,0.715283662,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_beach_C4,0.581357903,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_beach_C5,0.639033792,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_beach_C6,0.715201041,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_forest_C1,0.598821557,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_forest_C2,0.80534976,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_forest_C3,0.76745311,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_forest_C4,0.663004748,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_forest_C5,0.733410634,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_forest_C3,0.592214738,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_forest_C4,0.514185777,Familiar Object Familiar Background vs. Novel Object Novel Background
+o2_forest_C6,0.425907204,Familiar Object Familiar Background vs. Novel Object Novel Background
+o1_highway_C1,0.860614956,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_highway_C2,0.789680941,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_highway_C3,0.847120981,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_highway_C4,0.681576666,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_highway_C5,0.706469815,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_highway_C1,0.896121966,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_highway_C2,0.695178802,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_highway_C3,0.643366544,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_highway_C4,0.647303868,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_highway_C5,0.690925078,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_highway_C6,0.614870641,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_beach_C1,0.733962142,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_beach_C2,0.579727591,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_beach_C3,0.634215844,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_beach_C4,0.858648567,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_beach_C5,0.785277681,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_beach_C6,0.916735631,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_beach_C1,0.705731891,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_beach_C2,0.700869418,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_beach_C3,0.843615631,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_beach_C4,0.804880284,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_beach_C5,0.64967708,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_beach_C6,0.771270025,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_forest_C1,0.740266814,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_forest_C2,0.687969905,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_forest_C3,0.677778668,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_forest_C4,0.53959143,Familiar Object Novel Background vs. Novel Object Familiar Background
+o1_forest_C5,0.511136798,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_forest_C3,0.650524224,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_forest_C4,0.76172356,Familiar Object Novel Background vs. Novel Object Familiar Background
+o2_forest_C6,0.756944508,Familiar Object Novel Background vs. Novel Object Familiar Background
diff --git a/src/nett/analysis/NETT_merge_csvs.R b/src/nett/analysis/NETT_merge_csvs.R
new file mode 100644
index 00000000..c220e5c0
--- /dev/null
+++ b/src/nett/analysis/NETT_merge_csvs.R
@@ -0,0 +1,106 @@
+#!/usr/bin/env Rscript
+
+# merge_csvs.R
+# For a specified directory (see below), takes all of the csv files
+# and compiles them into a single data file
+
+# NOTE: For ease of use across many different experimental designs, 
+# this script assumes that all files use a common naming scheme with the 
+# following criteria:
+# 1) The agent ID number is the only number in the file name
+# 2) The filename ends with either train.csv or test.csv
+# for example "fork_side-agent3_train.csv"
+
+
+# Variables --------------------------------------------------------------------
+
+# Read in the user-specified variables:
+library(argparse)
+parser <- ArgumentParser(description="An executable R script for the Newborn Embodied Turing Tests to merge the log files across many agents")
+parser$add_argument("--logs-dir", type="character", dest="logs_dir",
+                    help="Working directory of the agents' log files",
+                    required=TRUE)
+parser$add_argument("--results-dir", type="character", dest="results_dir",
+                    help="Working directory to store the merged output",
+                    required=TRUE)
+parser$add_argument("--results-name", type="character", dest="results_name",
+                    help="File name for the R file storing the results",
+                    required=TRUE)
+parser$add_argument("--csv-train", type="character", dest="csv_train_name",
+                    help="File name for the csv file storing the training results",
+                    required=FALSE)
+parser$add_argument("--csv-test", type="character", dest="csv_test_name",
+                    help="File name for the csv file storing the testing results",
+                    required=FALSE)
+args <- parser$parse_args()
+data_wd <- args$logs_dir; results_wd <- args$results_dir; results_name <- args$results_name
+csv_train_name <- args$csv_train_name; csv_test_name <- args$csv_test_name
+
+# Set Zones:
+upper_x_lim <- 10
+lower_x_lim <- -10
+one_third <- (upper_x_lim - lower_x_lim)/3
+lower_bound <- lower_x_lim + one_third
+upper_bound <- upper_x_lim - one_third
+
+
+# Set Up -----------------------------------------------------------------------
+
+# Import libraries
+library(tidyverse)
+
+# Get all of the subdirectory csv filenames
+setwd(data_wd)
+train_files <- list.files(pattern="train.csv", recursive = TRUE)
+test_files <- list.files(pattern="test.csv", recursive = TRUE)
+
+# Main Function ----------------------------------------------------------------
+
+# This function reads in a single csv (later we'll lapply it across all files)
+read_data <- function(filename)
+{
+  # Read the csv file
+  data <- read.csv(filename)
+  
+  # Summarize by zones
+  data <- data %>%
+    mutate(left = case_when( agent.x < lower_bound ~ 1, agent.x >= lower_bound ~ 0)) %>%
+    mutate(right = case_when( agent.x > upper_bound ~ 1, agent.x <= upper_bound ~ 0)) %>%
+    mutate(middle = 1- left - right)
+  # Quick check to make sure that one and only one zone is chosen at each step
+  stopifnot(all( (data$left + data$right + data$middle == 1) ))
+  # Summarize at the episode level
+  data <- data %>%
+    group_by(Episode, left.monitor, right.monitor, correct.monitor, experiment.phase, imprint.cond, test.cond) %>%
+    summarise(.groups = "keep",
+      left_steps = sum(left), 
+      right_steps = sum(right), 
+      middle_steps = sum(middle)) %>%
+    mutate(Episode = as.numeric(Episode)) %>%
+    mutate(left.monitor = sub(" ", "", left.monitor)) %>%
+    mutate(right.monitor = sub(" ", "", right.monitor)) %>%
+    ungroup()
+  
+  # Add columns for original filename and agent ID number
+  data$filename <- basename(filename)
+  data$agent <- gsub("\\D", "", data$filename) # Only keep the number
+
+  return(data)
+}
+
+# Combine csv's and save results -----------------------------------------------
+
+# Combine all the training
+train_data <- lapply(train_files, FUN = read_data)
+train_data <- bind_rows(train_data)
+
+# Combine all the testing
+test_data <- lapply(test_files, FUN = read_data)
+test_data <- bind_rows(test_data)
+
+# Save it
+setwd(results_wd)
+save(train_data, test_data, file=results_name)
+if( !is.null(csv_train_name) ) write.csv(train_data, csv_train_name)
+if( !is.null(csv_test_name) ) write.csv(test_data, csv_test_name)
+
diff --git a/src/nett/analysis/NETT_train_viz.R b/src/nett/analysis/NETT_train_viz.R
new file mode 100644
index 00000000..22e3a98f
--- /dev/null
+++ b/src/nett/analysis/NETT_train_viz.R
@@ -0,0 +1,77 @@
+#!/usr/bin/env Rscript
+
+# NETT_train_viz.R
+
+# Before running this script, you need to run merge_csvs to merge all of the agents'
+# output into a single, standardized format dataframe for training and test data
+
+# Variables --------------------------------------------------------------------
+
+# Read in the user-specified variables:
+library(argparse)
+parser <- ArgumentParser(description="An executable R script for the Newborn Embodied Turing Tests to analyze test trials")
+parser$add_argument("--data-loc", type="character", dest="data_loc",
+                    help="Full filename (inc working directory) of the merged R data",
+                    required=TRUE)
+parser$add_argument("--results-wd", type="character", dest="results_wd",
+                    help="Working directory to save the resulting visualizations",
+                    required=TRUE)
+parser$add_argument("--ep-bucket", type="integer", dest="ep_bucket_size",
+                    help="How many episodes to group the x-axis by",
+                    required=TRUE)
+parser$add_argument("--num-episodes", type="integer", dest="num_episodes",
+                    help="How many episodes should be included",
+                    required=TRUE)
+args <- parser$parse_args()
+data_loc <- args$data_loc; results_wd <- args$results_wd; ep_bucket_size <- args$ep_bucket_size; num_episodes <- args$num_episodes
+
+# Set Up -----------------------------------------------------------------------
+
+library(tidyverse)
+
+load(data_loc)
+rm(test_data)
+setwd(results_wd)
+
+train_data_fixed <- train_data %>%
+  filter(Episode < num_episodes) %>%
+  # Create variables for correct/incorrect calculations
+  mutate(correct_steps = if_else(correct.monitor == " left", left_steps, right_steps)) %>%
+  mutate(incorrect_steps = if_else(correct.monitor == " left", right_steps, left_steps)) %>%
+  mutate(percent_correct = correct_steps / (correct_steps + incorrect_steps)) %>%
+  # Summarise data by condition, agent, and episode bucket for graphing
+  mutate(episode_block = Episode%/%ep_bucket_size + 1) %>%
+  group_by(imprint.cond, agent, episode_block) %>%
+  summarise(avgs = mean(percent_correct, na.rm = TRUE),
+            sd = sd(percent_correct, na.rm = TRUE),
+            count = length(percent_correct)) %>%
+  mutate(se = sd / sqrt(count)) %>%
+  # Convert numerical variables into correct type
+  mutate(episode_block = as.numeric(episode_block)) %>%
+  mutate(agent = as.numeric(agent)) %>%
+  ungroup()
+
+
+# Plot line graphs by imprinting condition -------------------------------------
+
+for (cond in unique(train_data_fixed$imprint.cond))
+{
+  data <- train_data_fixed %>%
+    filter(imprint.cond == cond)
+  
+  ggplot(data=data, aes(x=episode_block, y=avgs, color=as.factor(agent))) +
+    geom_line() +
+    theme_classic(base_size = 16) +
+    geom_hline(yintercept = .5, linetype = 2) +
+    xlab(sprintf("Groups of %d Episodes", ep_bucket_size)) + 
+    ylab("Average Time with Imprinted Object") +
+    scale_y_continuous(expand = c(0, 0), limits = c(0, 1), 
+                       breaks=seq(0,1,.1), labels = scales::percent) + 
+    scale_x_continuous(expand = c(0, 0), limits = c(0, num_episodes/ep_bucket_size), 
+                       breaks = seq(0, num_episodes / ep_bucket_size, 1)) +
+    theme(legend.position="none") 
+  
+  img_name <- paste0(cond, "_train.png")
+  ggsave(img_name)
+}
+
diff --git a/src/netts/body/__init__.py b/src/nett/body/__init__.py
similarity index 62%
rename from src/netts/body/__init__.py
rename to src/nett/body/__init__.py
index 73a91df9..c559c5c3 100644
--- a/src/netts/body/__init__.py
+++ b/src/nett/body/__init__.py
@@ -1,66 +1,75 @@
+"""
+Body module for nett
+
+.. module:: nett.body
+  :synopsis: Body module for nett
+
+.. moduleauthor:: Your Name <your.email@example.com>
+"""
+
 # all available body types
 types = ["basic", "two-eyed", "ragdoll"]
 
 # ASCII art
-ascii_basic = """                                                                   
-          ---------------::::::::::::::::::::::::::::::::::::::::::::::::::---------------          
-          ----------------::::::::::::::::::::::::::::::::::::::::::::::::----------------          
-          ------------------::::::::::::::::..::::::::::...::::::::::::::-----------------          
-          --------------------::::::::.:::::---------------::::.:::::::-------------------          
-          ----------------------:::::-------------:--------------::::---------------------          
-          ----------------------.------=====---:.  ..---========-----.--------------------          
-          --------------------:--====+++++======----======++****+====--.------------------          
-          ------------------:-===+***********==========*************+===------------------          
-          -----------------:==+****%%%%%%%%#***+++++++***%%%%%%%%%#***+==-:---------------          
-          -----------------=++**#%%%%%@@@%%%%**++++++**%%%@%%%%%@%%%%***+==:--------------          
-          ----------------=+***%@@@@@@@@@@@@@%**++++**%@@@@@@@@@@@@@@%***++=:-------------          
-          ---------------+++**@@@@@@@@@@@@@@@@********@@@@@@@@@@@@@@@@@**+++=-------------          
-          --------------++++*%@@@@@@@@@@@@@@@@+******+@@@@@@@@@@@@@@@@@*+++++=------------          
-          -------------=++*++@@@@@@@@@@@@@@@@++******++@@@@@@@@@@@@@@@@+++*++=:-----------          
-          -------------++**++%@@@@@@@@@@@@@@++********++@@@@@@@@@@@@@@@+++**++=-----------          
-          -----------.=+***+=+@@@@@@@@@@@@@++****:.****++@@@@@@@@@@@@@+==****+=:----------          
-          ------------+*****=-==@@@@@@@@===***-=+=+==.***===@@@@@@@@===-*****++-----------          
-          -----------=+*******+--=======****=+********-****+=======--=*******++-----------          
-          -----------=+*********************+**********:*********************++=----------          
-          -----------++********************-*##*****##*=**********************+=----------          
-          -----------************************###***###*************************+----------          
-          ====-------**######################################################**=---------=          
-          ========----*######################################################**------=====          
-          ==========--**#####################################################*----========          
-          ===========:.:*###################################################*=============          
-          =====::::......:################################################*.........======          
-          ===------**###+::::::....--###############################--.........+*---::.===          
-          ====.-===**#####********************------------======**********#####**-=--:--==          
-          ==+======**#########******************************************#######**====--===          
-          @*+++++=+**#################**************************###############*+=======+=          
-          @*++++++***##########################################################*+++++++++@          
-          ***+++*****##########################################################***++++++**          
-          +#*++******##########################################################*****++++**          
-          @##********#########################################################******+++**@          
-          @*###******#########################################################**********#@          
-          +*####****+#########################################################********###*          
-          ++######****#######################################################*+*****####*+          
-          %%#######***#######################################################****#######@@          
-          ++#######***#######################################################****#######++          
-          ##+######***######################################################*+**#######%@@          
-          ===#######**######################################################***########++*          
-          ====######***#####################################################***#######@@@%          
-          @@@@######***####################################################***########++++          
-          ===%%######***###################################################***#######====+          
-          =%@===######***#################################################***#######@@@@@@          
-          @@@@@@@######****##############################################***#######@%+====          
-          @========#####**+###########################################*+***######*==*%%===          
-          ==========*####*=*#########################################*@***######@@@@@@@@@@          
-          ============%*====*#######################################*==*######=========@%=          
-          @@@@%%%%%%%%%%%%%##*######################################+==@####============@%          
-          =======----%%##@@@@@#####################################*=-==@%===============@          
-          ==--------%@--------=###################################*@@@@@@@@@@@@@%%%%%%%%%%          
-          ---------%%----------=#################################=-------@%---=========**#          
-          --------%%-------------*#############################*----------@%------========          
-          %%%%%%%%%+--------------*###########################*------------%%---------====          
-          @@@@@@@@@@@@@@@@@@@@@@%%%%*#######################*--------------@%%-----------=          
-          ------%@------------------@@*###################*@%%%%%%%%%%%%%%%%@%------------          
-          -----%%-------------------%%--=*#############*%%----**@@@@@@@@@@@@@@@@@@@@@@@@@@          
-          ----%%-----------------===%@------=*******----%%-------------------@%%----------          
+ascii_basic = """
+          ---------------::::::::::::::::::::::::::::::::::::::::::::::::::---------------
+          ----------------::::::::::::::::::::::::::::::::::::::::::::::::----------------
+          ------------------::::::::::::::::..::::::::::...::::::::::::::-----------------
+          --------------------::::::::.:::::---------------::::.:::::::-------------------
+          ----------------------:::::-------------:--------------::::---------------------
+          ----------------------.------=====---:.  ..---========-----.--------------------
+          --------------------:--====+++++======----======++****+====--.------------------
+          ------------------:-===+***********==========*************+===------------------
+          -----------------:==+****%%%%%%%%#***+++++++***%%%%%%%%%#***+==-:---------------
+          -----------------=++**#%%%%%@@@%%%%**++++++**%%%@%%%%%@%%%%***+==:--------------
+          ----------------=+***%@@@@@@@@@@@@@%**++++**%@@@@@@@@@@@@@@%***++=:-------------
+          ---------------+++**@@@@@@@@@@@@@@@@********@@@@@@@@@@@@@@@@@**+++=-------------
+          --------------++++*%@@@@@@@@@@@@@@@@+******+@@@@@@@@@@@@@@@@@*+++++=------------
+          -------------=++*++@@@@@@@@@@@@@@@@++******++@@@@@@@@@@@@@@@@+++*++=:-----------
+          -------------++**++%@@@@@@@@@@@@@@++********++@@@@@@@@@@@@@@@+++**++=-----------
+          -----------.=+***+=+@@@@@@@@@@@@@++****:.****++@@@@@@@@@@@@@+==****+=:----------
+          ------------+*****=-==@@@@@@@@===***-=+=+==.***===@@@@@@@@===-*****++-----------
+          -----------=+*******+--=======****=+********-****+=======--=*******++-----------
+          -----------=+*********************+**********:*********************++=----------
+          -----------++********************-*##*****##*=**********************+=----------
+          -----------************************###***###*************************+----------
+          ====-------**######################################################**=---------=
+          ========----*######################################################**------=====
+          ==========--**#####################################################*----========
+          ===========:.:*###################################################*=============
+          =====::::......:################################################*.........======
+          ===------**###+::::::....--###############################--.........+*---::.===
+          ====.-===**#####********************------------======**********#####**-=--:--==
+          ==+======**#########******************************************#######**====--===
+          @*+++++=+**#################**************************###############*+=======+=
+          @*++++++***##########################################################*+++++++++@
+          ***+++*****##########################################################***++++++**
+          +#*++******##########################################################*****++++**
+          @##********#########################################################******+++**@
+          @*###******#########################################################**********#@
+          +*####****+#########################################################********###*
+          ++######****#######################################################*+*****####*+
+          %%#######***#######################################################****#######@@
+          ++#######***#######################################################****#######++
+          ##+######***######################################################*+**#######%@@
+          ===#######**######################################################***########++*
+          ====######***#####################################################***#######@@@%
+          @@@@######***####################################################***########++++
+          ===%%######***###################################################***#######====+
+          =%@===######***#################################################***#######@@@@@@
+          @@@@@@@######****##############################################***#######@%+====
+          @========#####**+###########################################*+***######*==*%%===
+          ==========*####*=*#########################################*@***######@@@@@@@@@@
+          ============%*====*#######################################*==*######=========@%=
+          @@@@%%%%%%%%%%%%%##*######################################+==@####============@%
+          =======----%%##@@@@@#####################################*=-==@%===============@
+          ==--------%@--------=###################################*@@@@@@@@@@@@@%%%%%%%%%%
+          ---------%%----------=#################################=-------@%---=========**#
+          --------%%-------------*#############################*----------@%------========
+          %%%%%%%%%+--------------*###########################*------------%%---------====
+          @@@@@@@@@@@@@@@@@@@@@@%%%%*#######################*--------------@%%-----------=
+          ------%@------------------@@*###################*@%%%%%%%%%%%%%%%%@%------------
+          -----%%-------------------%%--=*#############*%%----**@@@@@@@@@@@@@@@@@@@@@@@@@@
+          ----%%-----------------===%@------=*******----%%-------------------@%%----------
 """
-ascii_art = {'basic': ascii_basic}
\ No newline at end of file
+ascii_art = {'basic': ascii_basic}
diff --git a/src/nett/body/builder.py b/src/nett/body/builder.py
new file mode 100644
index 00000000..9ba6349f
--- /dev/null
+++ b/src/nett/body/builder.py
@@ -0,0 +1,142 @@
+"""The body of the agent in the environment."""
+from typing import Any
+from gym import Wrapper, Env
+
+from nett.body import types
+# from nett.body import ascii_art
+
+# this will have the necessary wrappers before the observations interact with the brain,
+# because it is the body that determines how the the observations will be processed.
+# specifically, in the case of the two-eyed agent, because the agent has two eyes, the observations are stereo
+# and need to be processed differently before they make it to the brain.
+# the body is the medium through which information travels from the environment to the brain.
+# the brain is limited by what the body can percieve and no information is objective.
+# NO INFORMATION IS OBJECTIVE (!!!!!!)
+class Body:
+    """Represents the body of an agent in an environment.
+
+    The body determines how observations from the environment are processed before they reach the brain.
+    It can apply wrappers to modify the observations and provide a different perception to the brain.
+
+    :param type: The type of the agent's body. Defaults to "basic".
+    :type type: str, optional
+    :param wrappers: List of wrappers to be applied to the environment. Defaults to None.
+    :type wrappers: list[Wrapper] | None, optional
+    :param dvs: Flag indicating whether the agent uses dynamic vision sensors. Defaults to False.
+    :type dvs: bool, optional
+
+    :raises ValueError: If the agent type is not valid.
+    :raises TypeError: If dvs is not a boolean.
+
+    Example:
+
+    >>> from nett import Body
+    >>> body = Body(type="basic", wrappers=None, dvs=False)
+    """
+
+    def __init__(self, type: str = "basic",
+                    wrappers: list[Any] | None = None,
+                    dvs: bool = False) -> None:
+        """
+        Constructor method
+        """
+        from nett import logger
+        self.logger = logger.getChild(__class__.__name__)
+        self.type = self._validate_agent_type(type)
+        self.wrappers = self._validate_wrappers(wrappers)
+        self.dvs = self._validate_dvs(dvs)
+
+    def _validate_agent_type(self, type: str) -> str:
+        """
+        Validate the agent type.
+
+        :param type: The type of the agent's body.
+        :type type: str
+
+        :returns: The validated agent type.
+        :rtype: str
+        """
+        if type not in types:
+            raise ValueError(f"agent type must be one of {types}")
+        return type
+
+
+    def _validate_dvs(self, dvs: bool) -> bool:
+        """
+        Validate the dvs flag.
+
+        :param dvs: The dvs flag.
+        :type dvs: bool
+
+        :returns: The validated dvs flag.
+        :rtype: bool
+
+        :raises TypeError: If dvs is not a boolean.
+        """
+        if not isinstance(dvs, bool):
+            raise TypeError("dvs should be a boolean [True, False]")
+        return dvs
+
+    def _validate_wrappers(self, wrappers: list[Any] | None) -> list[Any] | None:
+        """
+        Validate the wrappers.
+
+        :param wrappers: The list of wrappers.
+        :type wrappers: list[Wrapper] | None
+
+        :returns: The validated list of wrappers.
+        :rtype: list[Wrapper] | None
+
+        :raises ValueError: If any wrapper is not an instance of gym.Wrapper.
+        """
+        if wrappers is not None:
+            for wrapper in wrappers:
+                if not isinstance(wrapper, Wrapper):
+                    raise ValueError("Wrappers must inherit from gym.Wrapper")
+        return wrappers
+
+
+    def __call__(self, env: Env) -> Env:
+        """
+        Apply the registered wrappers to the environment.
+
+        :param env: The environment.
+        :type env: Env
+
+        :returns: The modified environment.
+        :rtype: Env
+        """
+        if self.wrappers:
+            for wrapper in self.wrappers:
+                env = wrapper(env)
+        return env
+
+    def __repr__(self) -> str:
+        """
+        Return a string representation of the Body object.
+
+        :returns: The string representation of the Body object.
+        :rtype: str
+        """
+        attrs = {k: v for k, v in vars(self).items() if k != "logger"}
+        return f"{self.__class__.__name__}({attrs!r})"
+
+
+    def __str__(self) -> str:
+        """
+        Return a string representation of the Body object.
+
+        :returns: The string representation of the Body object.
+        :rtype: str
+        """
+        attrs = {k: v for k, v in vars(self).items() if k != "logger"}
+        return f"{self.__class__.__name__}({attrs!r})"
+
+
+    def _register(self) -> None:
+        """
+        Register the body with the environment.
+
+        :raises NotImplementedError: This method is not implemented.
+        """
+        raise NotImplementedError
diff --git a/src/nett/brain/__init__.py b/src/nett/brain/__init__.py
new file mode 100644
index 00000000..9a7418dc
--- /dev/null
+++ b/src/nett/brain/__init__.py
@@ -0,0 +1,88 @@
+"""Initializes the brain module."""
+import ast
+from pathlib import Path
+
+import stable_baselines3
+import sb3_contrib
+from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
+
+from nett.brain import encoders
+
+def list_encoders() -> set[str]:
+    """
+    Return a set of all available encoders.
+
+    :return: Set of encoder names.
+    :rtype: set[str]
+    """
+    encoders_dir = Path.joinpath(Path(__file__).resolve().parent, 'encoders')
+    encoders = [encoder.stem for encoder in list(encoders_dir.iterdir()) if "__" not in str(encoder)]
+    # set is faster to access than a list
+    return set(encoders)
+
+encoders_list = list_encoders()
+
+def list_algorithms() -> set[str]:
+    """
+    Return a set of all available policy algorithms.
+
+    :return: Set of algorithm names.
+    :rtype: set[str]
+    """
+    sb3_policy_algorithms = [algorithm for algorithm in dir(stable_baselines3) if algorithm[0].isupper()]
+    sb3_contrib_policy_algorithms = [algorithm for algorithm in dir(sb3_contrib) if algorithm[0].isupper()]
+    available_policy_algorithms = sb3_policy_algorithms + sb3_contrib_policy_algorithms
+    # set is faster to access than a list
+    return set(available_policy_algorithms)
+
+algorithms = list_algorithms()
+
+# TODO (v0.3) return all available policy models programmatically
+def list_policies() -> set[str]:
+    """
+    Return a set of all available policy models.
+
+    :return: Set of policy model names.
+    :rtype: set[str]
+    """
+    return {'CnnPolicy', 'MlpPolicy', 'MultiInputPolicy', 'MultiInputLstmPolicy', 'CnnLstmPolicy'}
+
+policies = list_policies()
+
+# return encoder string to encoder class mapping
+# TODO (v0.3) optimized way to calculate and pass this dict around
+def get_encoder_dict():
+    """
+    Return a dictionary mapping encoder names to encoder class names.
+
+    :return: Dictionary mapping encoder names to encoder class names.
+    :rtype: dict[str, str]
+    """
+    encoders_dict = {}
+    encoders_dir = Path.joinpath(Path(__file__).resolve().parent, 'encoders')
+    # iterate through all files in the directory
+    for encoder_path in encoders_dir.iterdir():
+        if encoder_path.suffix == '.py' and "__" not in str(encoder_path):
+            module_name = encoder_path.stem
+            # read the source
+            with open(encoder_path) as f:
+                source = f.read()
+            # parse it
+            module = ast.parse(source)
+            # get the first class definition
+            encoder_class = [node for node in ast.walk(module) if isinstance(node, ast.ClassDef)][0]
+            # add to the dictionary
+            encoders_dict[module_name] = encoder_class.name
+    return encoders_dict
+
+# TODO (v0.3) return all available encoder classes programmatically
+# def get_encoder_dict():
+#     return {'cnnlstm': 'CNNLSTM',
+#             'cotracker': 'CoTracker',
+#             'dinov1': 'DinoV1',
+#             'dinov2': 'DinoV2',
+#             'vit': 'ViT',
+#             'resnet18': 'Resnet18CNN',
+#             'sam': 'SegmentAnything',
+#             'rnd': 'RND'}
+encoder_dict = get_encoder_dict()
diff --git a/src/nett/brain/builder.py b/src/nett/brain/builder.py
new file mode 100644
index 00000000..9d7b6745
--- /dev/null
+++ b/src/nett/brain/builder.py
@@ -0,0 +1,413 @@
+"""Module for the Brain class."""
+
+from typing import Any
+from pathlib import Path
+import inspect
+
+import torch
+import stable_baselines3
+import sb3_contrib
+import numpy as np
+import matplotlib.pyplot as plt
+from tqdm.auto import tqdm
+from stable_baselines3.common.callbacks import CallbackList, CheckpointCallback
+from sb3_contrib import RecurrentPPO
+from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
+from stable_baselines3.common.policies import BasePolicy
+from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm
+from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
+from stable_baselines3.common.env_util import make_vec_env
+from stable_baselines3.common.env_checker import check_env
+from stable_baselines3.common.monitor import Monitor
+from stable_baselines3.common.logger import configure
+from stable_baselines3.common import results_plotter
+
+from nett.brain import algorithms, policies, encoder_dict
+from nett.brain import encoders
+from nett.utils.callbacks import SupervisedSaveBestModelCallback, HParamCallback
+
+# TODO (v0.2): Extend with support for custom policy models
+# TODO (v0.2): should we move validation checks to utils under validations.py?
+class Brain:
+    """Represents the brain of an agent. 
+
+    The brain is made up of an encoder, policy, algorithm, reward function and the hyperparameters determined 
+    for these components such as the batch and buffer sizes. It produces a trained model based on the environment
+    data and the inputs received by the brain through the body.  
+    
+    :param encoder: The network used to extract features from the observations. Defaults to None.
+    :type encoder: Any | str, optional
+    :param embedding_dim: The dimension of the embedding space of the encoder. Defaults to None.
+    :type embedding_dim: int | None, optional
+    :param policy: The network used for defining the value and action networks. Defaults to None.
+    :type policy: Any | str | None, optional
+    :param algorithm: The optimization algorithm used for training the model. Defaults to None.
+    :type algorithm: Any | str | None, optional
+    :param reward: The type of reward used for training the brain. Defaults to "supervised".
+    :type reward: Any | str, optional
+    :param batch_size: The batch size used for training. Defaults to 512.
+    :type batch_size: int, optional
+    :param buffer_size: The buffer size used for training. Defaults to 2048.
+    :type buffer_size: int, optional
+    :param train_encoder: Whether to train the encoder or not. Defaults to False.
+    :type train_encoder: bool | None, optional
+    :param seed: The random seed used for training. Defaults to 12.
+    :type seed: int, optional
+
+    Example:
+
+    >>> from nett import Brain
+    >>> brain = Brain(policy='CnnPolicy', algorithm='PPO')
+    """
+
+    def __init__(self,
+                 encoder: Any | str = None,
+                 embedding_dim: int | None = None,
+                 policy: Any | str | None = None,
+                 algorithm: Any | str | None = None,
+                 reward: Any | str = "supervised",
+                 batch_size: int = 512,
+                 buffer_size: int = 2048,
+                 train_encoder: bool | None = False,
+                 seed: int = 12
+                 ) -> None:
+        """Constructor method
+        """
+        # Initialize logger
+        from nett import logger
+        self.logger = logger.getChild(__class__.__name__)
+        
+        # Set attributes
+        self.train_encoder = train_encoder
+        self.encoder = self._validate_encoder(encoder) if encoder else None
+        self.algorithm = self._validate_algorithm(algorithm) if algorithm else None
+        self.policy = self._validate_policy(policy) if policy else None
+        self.reward = self._validate_reward(reward) if reward else None
+        self.embedding_dim = embedding_dim
+        self.batch_size = batch_size
+        self.buffer_size = buffer_size
+        self.seed = seed
+
+    def train(self, env, iterations, device_type: str, device: int, paths: dict[str, Path]):
+        """
+        Train the brain.
+
+        :param env: The environment used for training.
+        :type env: Any
+        :param iterations: The number of training iterations.
+        :type iterations: int
+        :param device_type: The type of device used for training.
+        :type device_type: str
+        :param device: The device index used for training.
+        :type device: int
+        :param paths: The paths for saving logs, models, and plots.
+        :type paths: dict[str, Path]
+
+        :raises ValueError: If the environment fails the validation check.
+        
+        """
+        # validate environment
+        env = self._validate_env(env)
+
+        # initialize environment
+        log_path = paths['env_logs']
+        env = Monitor(env, str(log_path))
+        envs = make_vec_env(env_id=lambda : env, n_envs=1, seed=self.seed)
+
+        # build model
+        if self.encoder:
+            policy_kwargs = {
+                "features_extractor_class": self.encoder,
+                "features_extractor_kwargs": {
+                    "features_dim": inspect.signature(self.encoder).parameters['features_dim'].default
+                }
+            }
+        else:
+            policy_kwargs = {}
+        self.model = self.algorithm(self.policy,
+                                    envs,
+                                    batch_size=self.batch_size,
+                                    n_steps=self.buffer_size,
+                                    verbose=0,
+                                    policy_kwargs=policy_kwargs,
+                                    device=torch.device(device_type, device))
+
+        # setup tensorboard logger and attach to model
+        tb_logger = configure(str(paths['logs']), ["stdout", "csv", "tensorboard"])
+        self.model.set_logger(tb_logger)
+
+        # set encoder as eval only if train_encoder is not True
+        if not self.train_encoder:
+            self.model = self._set_encoder_as_eval(self.model)
+            self.logger.info(f'Encoder training is set to {str(self.train_encoder).upper()}')
+
+        # initialize callbacks
+        save_best_model_callback = SupervisedSaveBestModelCallback(
+            summary_freq=30000, 
+            save_dir=paths['model'], 
+            env_log_path=paths['env_logs'])
+        hparam_callback = HParamCallback()
+        checkpoint_callback = CheckpointCallback(save_freq=30000,
+                                                 save_path=paths["checkpoints"],
+                                                 name_prefix=self.algorithm.__name__,
+                                                 save_replay_buffer=True,
+                                                 save_vecnormalize=True)
+        callback_list = CallbackList([save_best_model_callback, hparam_callback, checkpoint_callback])
+
+        # train
+        self.logger.info(f"Total number of training steps: {iterations}")
+        self.model.learn(total_timesteps=iterations,
+                         tb_log_name=self.algorithm.__name__,
+                         progress_bar=True,
+                         callback=[callback_list])
+        self.logger.info("Training Complete")
+
+        # save
+        save_path = f"{paths['model'].joinpath('latest_model.zip')}"
+        self.save(save_path)
+        self.logger.info(f"Saved model at {save_path}")
+        # delete model to free up space
+        del self.model
+        # plot reward graph
+        self.plot_results(iterations=iterations,
+                          model_log_dir=paths['env_logs'],
+                          plots_dir=paths['plots'],
+                          name="reward_graph")
+
+    def test(self, env, iterations, model_path: str, record_prefix: str | None = None): # pylint: disable=unused-argument
+        """
+        Test the brain.
+
+        :param env: The environment used for testing.
+        :type env: Any
+        :param iterations: The number of testing iterations.
+        :type iterations: int
+        :param model_path: The path to the trained model.
+        :type model_path: str
+        :param record_prefix: The prefix for recording videos of the testing process, defaults to None.
+        :type record_prefix: str, optional
+        """
+        # load previously trained model from save_dir, if it exists
+        self.model = self.load(model_path)
+
+        # validate environment
+        env = self._validate_env(env)
+
+        # initialize environment
+        envs = make_vec_env(env_id=lambda : env, n_envs=1, seed=self.seed)
+
+        # for when algorithm is RecurrentPPO
+        if issubclass(self.algorithm, RecurrentPPO):
+            self.logger.info(f'Testing with {self.algorithm.__name__}')
+            self.logger.info(f"Total number of episodes: {iterations}")
+            num_envs = 1
+            for _ in tqdm(range(iterations)):
+                obs = env.reset()
+                # cell and hidden state of the LSTM
+                done, lstm_states = False, None
+                # episode start signals are used to reset the lstm states
+                episode_starts = np.ones((num_envs,), dtype=bool)
+                episode_length = 0
+                while not done:
+                    action, lstm_states = self.model.predict(obs, state=lstm_states, episode_start=episode_starts, deterministic=True)
+                    obs, _, done, _ = env.step(action) # obs, rewards, done, info
+                    episode_starts = done
+                    episode_length += 1
+                    env.render(mode="rgb_array")
+
+        # for all other algorithms
+        else:
+            self.logger.info(f'Testing with {self.algorithm.__name__}')
+            self.logger.info(f"Total number of testing steps: {iterations}")
+            obs = envs.reset()
+            for _ in tqdm(range(iterations)):
+                action, _ = self.model.predict(obs, deterministic=True) # action, states
+                obs, _, done, _ = envs.step(action) # obs, reward, done, info
+                if done:
+                    env.reset()
+                env.render(mode="rgb_array")
+
+    def save(self, path: str) -> None:
+        """
+        Save the trained model.
+
+        :param path: The path to save the model.
+        :type path: str
+        """
+        self.model.save(path)
+
+    def load(self, model_path: str | Path):
+        """
+        Load a trained model.
+
+        :param model_path: The path to the trained model.
+        :type model_path: str or Path
+
+        :return: The loaded model.
+        :rtype: Any
+        """
+        return self.algorithm.load(model_path)
+
+    def plot_results(self, iterations: int, model_log_dir: Path, plots_dir: Path, name: str) -> None:
+        """
+        Plot the training results.
+
+        :param iterations: The number of training iterations.
+        :type iterations: int
+        :param model_log_dir: The directory containing the training logs.
+        :type model_log_dir: Path
+        :param plots_dir: The directory to save the plots.
+        :type plots_dir: Path
+        :param name: The name of the plot.
+        :type name: str
+        """
+        results_plotter.plot_results([str(model_log_dir)],
+                                     iterations,
+                                     results_plotter.X_TIMESTEPS,
+                                     name)
+        plots_dir.mkdir(parents=True, exist_ok=True)
+        plt.savefig(plots_dir.joinpath(f"{name}.png"))
+        plt.clf()
+
+    def _validate_encoder(self, encoder: Any | str) -> BaseFeaturesExtractor:
+        """
+        Validate the encoder.
+
+        :param encoder: The encoder to validate.
+        :type encoder: Any or str
+
+        :return: The validated encoder.
+        :rtype: BaseFeaturesExtractor
+        """
+        # for when encoder is a string
+        if isinstance(encoder, str):
+            if encoder not in encoder_dict.keys():
+                raise ValueError(f"If a string, should be one of: {encoder_dict.keys()}")
+            encoder = getattr(encoders, encoder_dict[encoder])
+
+        # for when encoder is a custom PyTorch encoder
+        if isinstance(encoder, BaseFeaturesExtractor):
+            # TODO (v0.2) pass dummy torch.tensor on "meta" device to validate embedding dim
+            pass
+
+        if encoder and not hasattr(self, 'train_encoder'):
+            raise ValueError("encoder passed without setting train_encoder, should be one of: [True, False]")
+
+        return encoder
+
+    def _validate_algorithm(self, algorithm: Any | str) -> OnPolicyAlgorithm | OffPolicyAlgorithm:
+        """
+        Validate the optimization algorithm.
+
+        :param algorithm: The algorithm to validate.
+        :type algorithm: Any or str
+
+        :return: The validated algorithm.
+        :rtype: OnPolicyAlgorithm or OffPolicyAlgorithm
+        """
+        # for when policy is a string
+        if isinstance(algorithm, str):
+            if algorithm not in algorithms:
+                raise ValueError(f"If a string, should be one of: {algorithms}")
+            # check for the passed policy in stable_baselines3 as well as sb3-contrib
+            # at this point in the code, it is guaranteed to be in either of the two
+            try:
+                algorithm = getattr(stable_baselines3, algorithm)
+            except:
+                algorithm = getattr(sb3_contrib, algorithm)
+
+        # for when policy algorithm is custom
+        elif isinstance(algorithm, OnPolicyAlgorithm) or isinstance(algorithm, OffPolicyAlgorithm):
+            # TODO (v0.3) determine appropriate validation checks to be performed before passing
+            pass
+
+        else:
+            raise ValueError(f"Policy Algorithm should be either one of {algorithms} or a subclass of [{OnPolicyAlgorithm}, {OffPolicyAlgorithm}]")
+
+        return algorithm
+
+    def _validate_policy(self, policy: Any | str) -> str | BasePolicy:
+        """
+        Validate the policy model.
+
+        :param policy: The policy model to validate.
+        :type policy: Any or str
+
+        :return: The validated policy model.
+        :rtype: str or BasePolicy
+
+        :raises ValueError: If the policy is a string and not one of the supported policies.
+        :raises ValueError: If the policy is not a string or a subclass of BasePolicy.
+        """
+        # for when policy is a string
+        if isinstance(policy, str):
+            # support tested for PPO and RecurrentPPO only
+            if policy not in policies:
+                raise ValueError(f"If a string, should be one of: {policies}")
+
+        # for when policy is custom
+        elif isinstance(policy, BasePolicy):
+            # TODO (v0.3) determine appropriate validation checks to be performed before passing
+            pass
+
+        else:
+            raise ValueError(f"Policy Model should be either one of {policies} or a subclass of [{BasePolicy}]")
+
+        return policy
+
+    def _validate_reward(self, reward: Any | str) -> Any | str:
+        """
+        Validate the reward type.
+
+        :param reward: The reward type to validate.
+        :type reward: Any or str
+
+        :return: The validated reward type.
+        :rtype: Any or str
+
+        :raises ValueError: If the reward is a string and not one of the supported reward types.
+        """
+        # for when reward is a string
+        if isinstance(reward, str) and reward not in ['supervised', 'unsupervised']:
+            raise ValueError("If a string, should be one of: ['supervised', 'unsupervised']")
+        return reward
+
+    # TODO (v0.2) add typehinting for gym environments
+    def _validate_env(self, env) -> Any:
+        """
+        Validate the environment.
+
+        :param env: The environment to validate.
+        :type env: Any
+
+        :return: The validated environment.
+        :rtype: Any
+
+        :raises ValueError: If the environment fails the validation check.
+        """
+        try:
+            check_env(env)
+        except Exception as ex:
+            raise ValueError(f"Failed training env check with {str(ex)}")
+        return env
+
+    def _set_encoder_as_eval(self, model):
+        """
+        Set the encoder as evaluation mode and freeze its parameters.
+
+        :param model: The model containing the encoder.
+
+        :return: The model with the encoder set as evaluation mode.
+        """
+        model.policy.features_extractor.eval()
+        for param in model.policy.features_extractor.parameters():
+            param.requires_grad = False
+        return model
+
+    def __repr__(self) -> str:
+        attrs = {k: v for k, v in vars(self).items() if k != 'logger'}
+        return f"{self.__class__.__name__}({attrs!r})"
+
+    def __str__(self) -> str:
+        attrs = {k: v for k, v in vars(self).items() if k != 'logger'}
+        return f"{self.__class__.__name__}({attrs!r})"
diff --git a/src/netts/brain/encoders/__init__.py b/src/nett/brain/encoders/__init__.py
similarity index 72%
rename from src/netts/brain/encoders/__init__.py
rename to src/nett/brain/encoders/__init__.py
index 708986b8..06f7ff8e 100644
--- a/src/netts/brain/encoders/__init__.py
+++ b/src/nett/brain/encoders/__init__.py
@@ -1,9 +1,8 @@
-# simplify imports
+"""Simplifies imports for encoders"""
 from .resnet18 import Resnet18CNN
 from .dinov1 import DinoV1
 from .dinov2 import DinoV2
 from .sam import SegmentAnything
-from .rnd import RND
 from .cotracker import CoTracker
 from .vit import ViT
-from .cnnlstm import CNNLSTM
\ No newline at end of file
+from .cnnlstm import CNNLSTM
diff --git a/src/netts/brain/encoders/cnnlstm.py b/src/nett/brain/encoders/cnnlstm.py
similarity index 52%
rename from src/netts/brain/encoders/cnnlstm.py
rename to src/nett/brain/encoders/cnnlstm.py
index 4fa7024c..8a35d820 100644
--- a/src/netts/brain/encoders/cnnlstm.py
+++ b/src/nett/brain/encoders/cnnlstm.py
@@ -1,30 +1,25 @@
-#!/usr/bin/env python3
-
-import pdb
+"""CNNLSTM module for the brain"""
 import gym
-
-
 import torch as th
-import torch.nn as nn
-import torchvision
-
+from torch import nn
 from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
 
-
 class CNNLSTM(BaseFeaturesExtractor):
     """
-    :param observation_space: (gym.Space)
-    :param features_dim: (int) Number of features extracted.
-        This corresponds to the number of unit for the last layer.
+    CNNLSTM is a class that represents a convolutional neural network (CNN)
+    followed by a long short-term memory (LSTM) layer. It is used as a feature
+    extractor in reinforcement learning algorithms.
+
+    :param observation_space: The observation space of the environment.
+    :type observation_space: gym.Space
+    :param features_dim: Number of features extracted. This corresponds to the number of units for the last layer.
+    :type features_dim: int, optional, default 256
     """
-    def __init__(self, observation_space: gym.spaces.Box, 
-                  features_dim: int = 256):
-        rnn_hidden_size = 100
-        rnn_num_layers = 1
+    def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 256):
+        """Constructor method
+        """
         super(CNNLSTM, self).__init__(observation_space, features_dim)
-        
-        # We assume CxHxW images (channels first)
-        # Re-ordering will be done by pre-preprocessing or wrapper
+
         n_input_channels = observation_space.shape[0]
         self.cnn = nn.Sequential(
             nn.Conv2d(n_input_channels, 32, kernel_size=8, stride=4, padding=0),
@@ -40,53 +35,54 @@ def __init__(self, observation_space: gym.spaces.Box,
         with th.no_grad():
             n_flatten = self.cnn(th.as_tensor(observation_space.sample()[None]).float()).shape[1]
 
-         # define LSTM layer
+        # define LSTM layer
         hidden_size = 512
-        self.lstm = nn.LSTM(input_size = n_flatten, hidden_size = hidden_size,
-                            num_layers = 2, batch_first = True)
+        self.lstm = nn.LSTM(input_size=n_flatten, hidden_size=hidden_size,
+                            num_layers=2, batch_first=True)
+        
         # outputs
         self.linear = nn.Sequential(nn.Linear(hidden_size, features_dim), nn.ReLU())
-        
-        
-    def forward(self, observations: th.Tensor) :
+
+    def forward(self, observations: th.Tensor):
         """
-        observations = observations.unsqueeze(0)
-        
-        batch_size, seq_length, c, h, w = observations.shape
-        ii = 0
-        y = self.cnn((observations[:,ii]))
-        
-        #outputs = [sen_len, batch_size, hid_dim * n_directions]
-        out, (hn, cn) = self.lstm(y.unsqueeze(1))
-        out = self.linear(out[:,-1]) 
-        
-        Args:
-            observations (th.Tensor): _description_
+        Forward pass of the CNNLSTM.
 
-        Returns:
-            _type_: _description_
+        :param observations: The input observations.
+        :type observations: th.Tensor
+        :return: The extracted features.
+        :rtype: th.Tensor
         """
         x = observations # original shape -> (length, batchsize, obs_size)
-        T,B, *_ = x.shape
-        
+        # T,B, *_ = x.shape
+
         # Pass through CNN layers
         x = self.cnn(x)
-        
+
         # Flatten the output for LSTM
         x = x.view(x.size(0), x.size(1), -1)
-        
+
         # Pass through LSTM layer
         x, _ = self.lstm(x)
-        
+
         # Get the last time step's output and apply the fully connected layer
         x = self.linear(x[:, -1, :])
-        
-        return x 
 
-    
+        return x
 
 class Identity(nn.Module):
+    """Identity module
+    
+    This module is used to return the input tensor as is.
+    
+    :param nn.Module: PyTorch module
+    
+    :return: Identity module
+    :rtype: nn.Module
+    """
     def __init__(self):
+        """Constructor method"""
         super(Identity, self).__init__()
+
     def forward(self, x):
-        return x
\ No newline at end of file
+        """Forward pass"""
+        return x
diff --git a/src/netts/brain/encoders/cotracker.py b/src/nett/brain/encoders/cotracker.py
similarity index 66%
rename from src/netts/brain/encoders/cotracker.py
rename to src/nett/brain/encoders/cotracker.py
index af61c2b5..9ffb13b7 100644
--- a/src/netts/brain/encoders/cotracker.py
+++ b/src/nett/brain/encoders/cotracker.py
@@ -1,63 +1,64 @@
 #!/usr/bin/env python3
-
-
-import pdb
+"""CoTracker feature extractor"""
+# import pdb
 import gym
-
-
 import torch as th
-import torch.nn as nn
-import torchvision
-import timm
-import numpy as np
 from torchvision.transforms import Compose
 from torchvision.transforms import Resize, CenterCrop, Normalize, InterpolationMode
 from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
-from collections import OrderedDict
+# from collections import OrderedDict
 
 class CoTracker(BaseFeaturesExtractor):
     """
-    :param observation_space: (gym.Space)
-    :param features_dim: (int) Number of features extracted.
-        This corresponds to the number of unit for the last layer.
-    """
+    Initialize the CoTracker feature extractor.
 
+    :param observation_space: The observation space of the environment.
+    :type observation_space: gym.spaces.Box
+    :param features_dim: Number of features extracted. This corresponds to the number of units for the last layer.
+    :type features_dim: int
+    """
     def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 384):
+        """Constructor method
+        """
         super(CoTracker, self).__init__(observation_space, features_dim)
         self.n_input_channels = observation_space.shape[0]
-        
-        
-        self.transforms = Compose([Resize(size=256, 
-                                          interpolation=InterpolationMode.BICUBIC, 
-                                          max_size=None, 
+
+        self.transforms = Compose([Resize(size=256,
+                                          interpolation=InterpolationMode.BICUBIC,
+                                          max_size=None,
                                           antialias=True),
                                    CenterCrop(size=(224, 224)),
-                                   Normalize(mean=th.tensor([0.485, 0.456, 0.406]), 
+                                   Normalize(mean=th.tensor([0.485, 0.456, 0.406]),
                                              std=th.tensor([0.229, 0.224, 0.225]))])
-        
-        
+
         n_input_channels = observation_space.shape[0]
         print("N_input_channels", n_input_channels)
-        
+
         model = th.hub.load("facebookresearch/co-tracker", "cotracker_w8")
         modules = []
         for name, layer in model.named_modules():
-            if name == 'model.fnet':
+            if name == "model.fnet":
                 modules.append(layer)
-                
-        self.cnn = (th.nn.Sequential(*modules))
+
+        self.cnn = th.nn.Sequential(*modules)
         self.linear = th.nn.Sequential(th.nn.Linear(128*16*16, 512),
                                             th.nn.ReLU())
-        
+
         #dummy_x = th.zeros((1, 1, 3, 64, 64))
         #conv_out_size = (self.model(dummy_x))
         #pdb.set_trace()
         ## remove classifier
         #self.model = th.nn.Sequential(OrderedDict([*(list(model.named_children())[:-1])]))
-        
-        
-        
+
     def forward(self, observations: th.Tensor) -> th.Tensor:
+        """
+        Forward pass of the CoTracker feature extractor.
+
+        :param observations: The input observations.
+        :type observations: th.Tensor
+        :return: The extracted features.
+        :rtype: th.Tensor
+        """
         # Cut off image
         # reshape to from vector to W*H
         # gray to color transform
@@ -66,4 +67,4 @@ def forward(self, observations: th.Tensor) -> th.Tensor:
         # return
         x = self.cnn(observations)
         x = x.view(x.size(0), -1)
-        return self.linear(x)
\ No newline at end of file
+        return self.linear(x)
diff --git a/src/netts/brain/encoders/dinov1.py b/src/nett/brain/encoders/dinov1.py
similarity index 52%
rename from src/netts/brain/encoders/dinov1.py
rename to src/nett/brain/encoders/dinov1.py
index 27d0e720..b490bcbf 100644
--- a/src/netts/brain/encoders/dinov1.py
+++ b/src/nett/brain/encoders/dinov1.py
@@ -1,3 +1,4 @@
+"""DINO (Emerging Properties in Self-Supervised Vision Transformers) model"""
 import gym
 import torch
 import timm
@@ -5,27 +6,39 @@
 from torchvision.transforms import Compose, Resize, CenterCrop, Normalize, InterpolationMode
 from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
 
-
 class DinoV1(BaseFeaturesExtractor):
     """
-    :param observation_space: (gym.Space)
-    :param features_dim: (int) Number of features extracted.
-        This corresponds to the number of unit for the last layer.
+    Initialize DinoV1 feature extractor.
+
+    :param observation_space: The observation space of the environment.
+    :type observation_space: gym.spaces.Box
+    :param features_dim: Number of features extracted. This corresponds to the number of units for the last layer.
+    :type features_dim: int
     """
     def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 384):
+        """Constructor method
+        """
         super(DinoV1, self).__init__(observation_space, features_dim)
         self.n_input_channels = observation_space.shape[0]
-        self.transforms = Compose([Resize(size=248, 
-                                          interpolation=InterpolationMode.BICUBIC, 
-                                          max_size=None, 
+        self.transforms = Compose([Resize(size=248,
+                                          interpolation=InterpolationMode.BICUBIC,
+                                          max_size=None,
                                           antialias=True),
                                    CenterCrop(size=(224, 224)),
-                                   Normalize(mean=torch.tensor([0.4850, 0.4560, 0.4060]), 
-                                             std=torch.tensor([0.2290, 0.2240, 0.2250]))])        
-        self.model = timm.create_model('vit_small_patch8_224.dino', 
-                                       in_chans=self.n_input_channels, 
-                                       num_classes=0, 
+                                   Normalize(mean=torch.tensor([0.4850, 0.4560, 0.4060]),
+                                             std=torch.tensor([0.2290, 0.2240, 0.2250]))])
+        self.model = timm.create_model('vit_small_patch8_224.dino',
+                                       in_chans=self.n_input_channels,
+                                       num_classes=0,
                                        pretrained=True)
-        
+
     def forward(self, observations: torch.Tensor) -> torch.Tensor:
-        return self.model(self.transforms(observations))
\ No newline at end of file
+        """
+        Forward pass of the DinoV1 model.
+
+        :param observations: The input observations.
+        :type observations: torch.Tensor
+        :return: The extracted features.
+        :rtype: torch.Tensor
+        """
+        return self.model(self.transforms(observations))
diff --git a/src/nett/brain/encoders/dinov2.py b/src/nett/brain/encoders/dinov2.py
new file mode 100644
index 00000000..2c32f73f
--- /dev/null
+++ b/src/nett/brain/encoders/dinov2.py
@@ -0,0 +1,60 @@
+"""DinoV2 feature extractor
+
+This module provides a feature extractor based on the DINOv2 model. It takes in observations from an environment and extracts features using the DINOv2 model.
+
+Example:
+
+    >>> observation_space = gym.spaces.Box(low=0, high=255, shape=(3, 84, 84), dtype=np.uint8)
+    >>> features_dim = 384
+    >>> extractor = DinoV2(observation_space, features_dim)
+    >>> observations = torch.randn(1, 3, 84, 84)
+    >>> features = extractor.forward(observations)
+
+Attributes:
+    n_input_channels (int): The number of input channels in the observation space.
+    transforms (torchvision.transforms.Compose): A series of image transformations applied to the observations.
+    model (torch.nn.Module): The DINOv2 model used for feature extraction.
+
+Methods:
+    __init__(observation_space, features_dim): Initializes the DinoV2 feature extractor.
+    forward(observations): Performs a forward pass of the DinoV2 feature extractor.
+
+"""
+
+import gym
+import torch
+
+from torchvision.transforms import Compose, Resize, CenterCrop, Normalize, InterpolationMode
+from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
+
+class DinoV2(BaseFeaturesExtractor):
+    """
+    DinoV2 is a feature extractor based on the DINOv2 model.
+
+    :param observation_space: The observation space of the environment. 
+    :type observation_space: gym.spaces.Box
+    :param features_dim: Number of features extracted. This corresponds to the number of units for the last layer. Defaults to 384.
+    :type features_dim: int, optional
+    """
+    def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 384):
+        super(DinoV2, self).__init__(observation_space, features_dim)
+        """Constructor method
+        """
+        self.n_input_channels = observation_space.shape[0]
+        self.transforms = Compose([Resize(size=256,
+                                          interpolation=InterpolationMode.BICUBIC,
+                                          max_size=None,
+                                          antialias=True),
+                                   CenterCrop(size=(224, 224)),
+                                   Normalize(mean=torch.tensor([0.485, 0.456, 0.406]),
+                                             std=torch.tensor([0.229, 0.224, 0.225]))])
+        self.model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14", pretrained=True)
+
+    def forward(self, observations: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the DinoV2 feature extractor.
+
+        :param observations: (torch.Tensor) The input observations.
+        :return: (torch.Tensor) The extracted features.
+        """
+        return self.model(self.transforms(observations))
diff --git a/src/netts/brain/encoders/resnet18.py b/src/nett/brain/encoders/resnet18.py
similarity index 58%
rename from src/netts/brain/encoders/resnet18.py
rename to src/nett/brain/encoders/resnet18.py
index 9bea45af..32d75fe6 100644
--- a/src/netts/brain/encoders/resnet18.py
+++ b/src/nett/brain/encoders/resnet18.py
@@ -1,20 +1,45 @@
+"""
+This module defines the Resnet18CNN class, which is a custom feature extractor
+based on the ResNet-18 architecture. It is used for encoding observations in
+the Newborn Embodied Turing Test project.
+
+The Resnet18CNN class inherits from the BaseFeaturesExtractor class provided
+by the stable_baselines3 library. It takes an observation space and the desired
+number of features as input and extracts features using the ResNet-18 model.
+
+The ResNet-18 architecture consists of several residual blocks, each containing
+two convolutional layers and a skip connection. The final features are obtained
+by applying a linear layer to the output of the last residual block.
+
+The ResBlock class and the ResNet_18 class are helper classes used by the
+Resnet18CNN class to define the residual blocks and the overall ResNet-18
+architecture, respectively.
+
+Example usage:
+
+    observation_space = gym.spaces.Box(low=0, high=255, shape=(3, 84, 84), dtype=np.uint8)
+    features_dim = 256
+    encoder = Resnet18CNN(observation_space, features_dim)
+    features = encoder(observation)
+
+"""
+
 #!/usr/bin/env python3
 
-import pdb
+# import pdb
 import gym
 
-
 import torch as th
-import torch.nn as nn
-import torchvision
+from torch import nn
 
 from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
 
 class Resnet18CNN(BaseFeaturesExtractor):
     """
-    :param observation_space: (gym.Space)
-    :param features_dim: (int) Number of features extracted.
-        This corresponds to the number of unit for the last layer.
+    Custom feature extractor based on the ResNet-18 architecture.
+
+    :param observation_space: (gym.Space) The observation space of the environment.
+    :param features_dim: (int) Number of features to be extracted.
     """
 
     def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 256):
@@ -30,10 +55,16 @@ def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 256):
             n_flatten = self.cnn(
                 th.as_tensor(observation_space.sample()[None]).float()
             ).shape[1]
-        
+
         self.linear = nn.Sequential(nn.Linear(n_flatten, features_dim), nn.ReLU())
-        
+
     def forward(self, observations: th.Tensor) -> th.Tensor:
+        """
+        Forward pass of the feature extractor.
+
+        :param observations: (torch.Tensor) The input observations.
+        :return: (torch.Tensor) The extracted features.
+        """
         # Cut off image
         # reshape to from vector to W*H
         # gray to color transform
@@ -43,7 +74,11 @@ def forward(self, observations: th.Tensor) -> th.Tensor:
         return self.linear(self.cnn(observations))
 
 ## reference - online
-class ResBlock(nn.Module): 
+class ResBlock(nn.Module):
+    """
+    Residual block used in the ResNet-18 architecture.
+    """
+
     def __init__(self, in_channels, out_channels, identity_downsample=None, stride=1):
         super(ResBlock, self).__init__()
         self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
@@ -52,8 +87,14 @@ def __init__(self, in_channels, out_channels, identity_downsample=None, stride=1
         self.bn2 = nn.BatchNorm2d(out_channels)
         self.relu = nn.ReLU()
         self.identity_downsample = identity_downsample
-        
+
     def forward(self, x):
+        """
+        Forward pass of the residual block.
+
+        :param x: (torch.Tensor) The input tensor.
+        :return: (torch.Tensor) The output tensor.
+        """
         identity = x
         x = self.conv1(x)
         x = self.bn1(x)
@@ -67,56 +108,76 @@ def forward(self, x):
         return x
 
 class ResNet_18(nn.Module):
-    
+    """
+    ResNet-18 architecture used in the Resnet18CNN class.
+    """
+
     def __init__(self, image_channels, num_classes):
-        
         super(ResNet_18, self).__init__()
         self.in_channels = 64
         self.conv1 = nn.Conv2d(image_channels, 64, kernel_size=7, stride=2, padding=3)
         self.bn1 = nn.BatchNorm2d(64)
         self.relu = nn.ReLU()
         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
-        
+
         #resnet layers
         self.layer1 = self.__make_layer(64, 64, stride=1)
         self.layer2 = self.__make_layer(64, 128, stride=2)
         self.layer3 = self.__make_layer(128, 256, stride=2)
         self.layer4 = self.__make_layer(256, 512, stride=2)
-        
+
         self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
         self.fc = nn.Linear(512, num_classes)
-        
+
     def __make_layer(self, in_channels, out_channels, stride):
-        
+        """
+        Helper function to create a residual layer.
+
+        :param in_channels: (int) Number of input channels.
+        :param out_channels: (int) Number of output channels.
+        :param stride: (int) Stride of the convolutional layers.
+        :return: (nn.Sequential) The residual layer.
+        """
         identity_downsample = None
         if stride != 1:
             identity_downsample = self.identity_downsample(in_channels, out_channels)
-            
+
         return nn.Sequential(
-            ResBlock(in_channels, out_channels, identity_downsample=identity_downsample, stride=stride), 
+            ResBlock(in_channels, out_channels, identity_downsample=identity_downsample, stride=stride),
             ResBlock(out_channels, out_channels)
         )
-        
+
     def forward(self, x):
-        
+        """
+        Forward pass of the ResNet-18 architecture.
+
+        :param x: (torch.Tensor) The input tensor.
+        :return: (torch.Tensor) The output tensor.
+        """
         x = self.conv1(x)
         x = self.bn1(x)
         x = self.relu(x)
         x = self.maxpool(x)
-        
+
         x = self.layer1(x)
         x = self.layer2(x)
         x = self.layer3(x)
         x = self.layer4(x)
-        
+
         x = self.avgpool(x)
         x = x.view(x.shape[0], -1)
         x = self.fc(x)
-        return x 
-    
+        return x
+
     def identity_downsample(self, in_channels, out_channels):
-        
+        """
+        Helper function to create an identity downsample layer.
+
+        :param in_channels: (int) Number of input channels.
+        :param out_channels: (int) Number of output channels.
+        :return: (nn.Sequential) The identity downsample layer.
+        """
         return nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1), 
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1),
             nn.BatchNorm2d(out_channels)
-        )
\ No newline at end of file
+        )
diff --git a/src/netts/brain/encoders/sam.py b/src/nett/brain/encoders/sam.py
similarity index 55%
rename from src/netts/brain/encoders/sam.py
rename to src/nett/brain/encoders/sam.py
index fbe40f4b..8ea9b15a 100644
--- a/src/netts/brain/encoders/sam.py
+++ b/src/nett/brain/encoders/sam.py
@@ -1,11 +1,12 @@
 #!/usr/bin/env python3
-import pdb
+"""
+This module contains the implementation of the SegmentAnything class, which is a custom feature extractor
+for image observations in a gym environment. It uses a pre-trained model from the timm library to extract
+features from the input images.
+"""
 import gym
 
-
 import torch as th
-import torch.nn as nn
-import torchvision
 import timm
 from torchvision.transforms import Compose
 from torchvision.transforms import Resize, CenterCrop, Normalize, InterpolationMode
@@ -14,31 +15,36 @@
 
 class SegmentAnything(BaseFeaturesExtractor):
     """
-    :param observation_space: (gym.Space)
-    :param features_dim: (int) Number of features extracted.
-        This corresponds to the number of unit for the last layer.
+    Custom feature extractor for image observations in a gym environment.
+
+    :param observation_space: (gym.Space) The observation space of the environment.
+    :param features_dim: (int) Number of features extracted. This corresponds to the number of units for the last layer.
     """
 
     def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 384):
         super(SegmentAnything, self).__init__(observation_space, features_dim)
         self.n_input_channels = observation_space.shape[0]
-        self.transforms = Compose([Resize(size=256, 
-                                          interpolation=InterpolationMode.BICUBIC, 
-                                          max_size=None, 
+        self.transforms = Compose([Resize(size=256,
+                                          interpolation=InterpolationMode.BICUBIC,
+                                          max_size=None,
                                           antialias=True),
                                    CenterCrop(size=(224, 224)),
-                                   Normalize(mean=th.tensor([0.485, 0.456, 0.406]), 
+                                   Normalize(mean=th.tensor([0.485, 0.456, 0.406]),
                                              std=th.tensor([0.229, 0.224, 0.225]))])
-        
+
         n_input_channels = observation_space.shape[0]
         print("N_input_channels", n_input_channels)
-        
-        
-        self.model = timm.create_model('samvit_base_patch16.sa1b',pretrained=True,
-                                  num_classes=0)  # remove classifier nn.Linear)
-        
-        
+
+        self.model = timm.create_model("samvit_base_patch16.sa1b", pretrained=True,
+                                       num_classes=0)  # remove classifier th.nn.Linear)
+
     def forward(self, observations: th.Tensor) -> th.Tensor:
+        """
+        Forward pass of the feature extractor.
+
+        :param observations: (th.Tensor) The input observations.
+        :return: (th.Tensor) The extracted features.
+        """
         # Cut off image
         # reshape to from vector to W*H
         # gray to color transform
@@ -46,4 +52,3 @@ def forward(self, observations: th.Tensor) -> th.Tensor:
         # Concat features to the rest of observation vector
         # return
         return self.model(self.transforms(observations))
-
diff --git a/src/nett/brain/encoders/vit.py b/src/nett/brain/encoders/vit.py
new file mode 100644
index 00000000..bf97ebce
--- /dev/null
+++ b/src/nett/brain/encoders/vit.py
@@ -0,0 +1,51 @@
+"""ViT (Vision Transformer) encoder"""
+import gym
+import torch
+import timm
+
+from torchvision.transforms import Compose
+from torchvision.transforms import Resize, CenterCrop, Normalize, InterpolationMode
+from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
+
+class ViT(BaseFeaturesExtractor):
+    """
+    ViT is a feature extractor based on the Vision Transformer model.
+
+    Args:
+        observation_space (gym.spaces.Box): The observation space of the environment.
+        features_dim (int, optional): The dimension of the extracted features. Defaults to 384.
+    """
+    def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 384):
+        """
+        Initializes the ViT (Vision Transformer) encoder.
+
+        Args:
+            observation_space (gym.spaces.Box): The observation space of the environment.
+            features_dim (int, optional): The dimension of the extracted features. Defaults to 384.
+        """
+        super(ViT, self).__init__(observation_space, features_dim)
+        self.n_input_channels = observation_space.shape[0]
+        self.transforms = Compose([Resize(size=248,
+                                          interpolation=InterpolationMode.BICUBIC,
+                                          max_size=None,
+                                          antialias=True),
+                                   CenterCrop(size=(224, 224)),
+                                   Normalize(mean=torch.tensor([0.4850, 0.4560, 0.4060]),
+                                             std=torch.tensor([0.2290, 0.2240, 0.2250]))])
+
+        self.model = timm.create_model("vit_small_patch8_224.dino",
+                                       in_chans=self.n_input_channels,
+                                       num_classes=0,
+                                       pretrained=False)
+
+    def forward(self, observations: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the ViT encoder.
+
+        Args:
+            observations (torch.Tensor): The input observations.
+
+        Returns:
+            torch.Tensor: The extracted features.
+        """
+        return self.model(self.transforms(observations))
diff --git a/src/nett/environment/__init__.py b/src/nett/environment/__init__.py
new file mode 100644
index 00000000..13f4e64b
--- /dev/null
+++ b/src/nett/environment/__init__.py
@@ -0,0 +1,6 @@
+"""
+Initializes the environment module.
+"""
+# simplify imports
+from nett.environment.configs import list_configs
+from nett.environment.configs import IdentityAndView, Parsing, Slowness, Smoothness, Binding, OneShotViewInvariant
diff --git a/src/nett/environment/builder.py b/src/nett/environment/builder.py
new file mode 100644
index 00000000..217ebf2e
--- /dev/null
+++ b/src/nett/environment/builder.py
@@ -0,0 +1,238 @@
+"""Module for the Environment class."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from typing import Optional
+
+import numpy as np
+from gym import Wrapper
+from mlagents_envs.environment import UnityEnvironment
+
+# checks to see if ml-agents tmp files have the proper permissions
+try :
+    from mlagents_envs.envs.unity_gym_env import UnityToGymWrapper
+except PermissionError as _:
+     raise PermissionError("Directory '/tmp/ml-agents-binaries' is not accessible. Please change permissions of the directory and its subdirectories ('tmp' and 'binaries') to 1777 or delete the entire directory and try again.")
+
+from nett.environment import configs
+from nett.environment.configs import NETTConfig, list_configs
+from nett.utils.environment import Logger, port_in_use
+
+class Environment(Wrapper):
+    """
+    Represents the environment where the agent lives.
+
+    The environment is the source of all input data streams to train the brain of the agent. 
+    It accepts a Unity Executable and wraps it around as a Gym environment by leveraging the UnityEnvironment 
+    class from the mlagents_envs library.
+
+    It provides a convenient interface for interacting with the Unity environment and includes methods for initializing the environment, rendering frames, taking steps, resetting the environment, and logging messages.
+
+    :param config: The configuration for the environment. It can be either a string representing the name of a pre-defined configuration, or an instance of the NETTConfig class.
+    :type config: str | NETTConfig
+    :param executable_path: The path to the Unity executable file.
+    :type executable_path: str
+    :param display: The display number to use for the Unity environment. Defaults to 0.
+    :type display: int, optional
+    :param base_port: The base port number to use for communication with the Unity environment. Defaults to 5004.
+    :type base_port: int, optional
+    :param record_chamber: Whether to record the chamber. Defaults to False.
+    :type record_chamber: bool, optional
+    :param record_agent: Whether to record the agent. Defaults to False.
+    :type record_agent: bool, optional
+    :param recording_frames: The number of frames to record. Defaults to 1000.
+    :type recording_frames: int, optional
+
+    :raises ValueError: If the configuration is not a valid string or an instance of NETTConfig.
+
+    Example:
+
+    >>> from nett import Environment
+    >>> env = Environment(config="identityandview", executable_path="path/to/executable")
+    """
+    def __init__(self,
+                 config: str | NETTConfig,
+                 executable_path: str,
+                 display: int = 0,
+                 base_port: int = 5004,
+                 record_chamber: bool = False,
+                 record_agent: bool = False,
+                 recording_frames: int = 1000) -> None:
+        """Constructor method
+        """
+        from nett import logger
+        self.logger = logger.getChild(__class__.__name__)
+        self.config = self._validate_config(config)
+        # TODO (v0.4) what might be a way to check if it is a valid executable path?
+        self.executable_path = executable_path
+        self.base_port = base_port
+        self.record_chamber = record_chamber
+        self.record_agent = record_agent
+        self.recording_frames = recording_frames
+        self.display = display
+
+        # set the correct permissions on the executable
+        self._set_executable_permission()
+        # set the display for Unity environment
+        self._set_display()
+
+    def _validate_config(self, config: str | NETTConfig) -> NETTConfig:
+        """
+        Validates the configuration for the environment.
+
+        :param config: The configuration to validate.
+        :type config: str | NETTConfig
+        :return: The validated configuration.
+        :rtype: NETTConfig
+        :raises ValueError: If the configuration is not a valid string or an instance of NETTConfig.
+        """
+        # for when config is a str
+        if isinstance(config, str):
+            config_dict = {config_str.lower(): config_str for config_str in list_configs()}
+            if config not in config_dict.keys():
+                raise ValueError(f"Should be one of {config_dict.keys()}")
+
+            config = getattr(configs, config_dict[config])()
+
+        # for when config is a NETTConfig
+        elif isinstance(config, NETTConfig):
+            pass
+
+        else:
+            raise ValueError(f"Should either be one of {list(config_dict.keys())} or a subclass of NETTConfig")
+
+        return config
+
+    def _set_executable_permission(self) -> None:
+        """
+        Sets the executable permission for the Unity executable file.
+
+        :return: None
+        """
+        subprocess.run(["chmod", "-R", "755", self.executable_path], check=True)
+        self.logger.info("Executable permission is set")
+
+    def _set_display(self) -> None:
+        """
+        Sets the display environment variable for the Unity environment.
+
+        :return: None
+        """
+        os.environ["DISPLAY"] = str(f":{self.display}")
+        self.logger.info("Display is set")
+
+    
+    # copied from __init__() of chickai_env_wrapper.py (legacy)
+    # TODO (v0.3) Critical refactor, don"t like how this works, extremely error prone.
+    # how can we build + constraint arguments better? something like an ArgumentParser sounds neat
+    # TODO (v0.3) fix random_pos logic inside of Unity code
+    def initialize(self, mode: str, **kwargs) -> Environment:
+        """
+        Initializes the environment with the given mode and arguments.
+
+        :param mode: The mode to set the environment for training or testing or both. 
+        :type mode: str
+        :param kwargs: The arguments to pass to the environment.
+        :type kwargs: Any
+        :return: The initialized environment.
+        :rtype: Environment
+        """
+        args = []
+
+        # from environment arguments
+        if self.recording_frames:
+            args.extend(["--recording-steps", str(self.recording_frames)])
+        if self.record_chamber:
+            args.extend(["--record-chamber", "true"])
+        if self.record_agent:
+            args.extend(["--record-agent", "true"])
+
+        # from runtime
+        args.extend(["--mode", f"{mode}-{kwargs['condition']}"])
+        if kwargs.get("rec_path", None):
+            args.extend(["--log-dir", f"{kwargs['rec_path']}/"])
+        # needs to fixed in Unity code where the default is always false
+        if mode == "train":
+            args.extend(["--random-pos", "true"])
+        if kwargs.get("rewarded", False):
+            args.extend(["--rewarded", "true"])
+        # TODO: Discuss this with Manju, may be a MAJOR bug
+        self.step_per_episode = kwargs.get("episode_steps", 200)
+        if kwargs.get("episode_steps", False):
+            args.extend(["--episode-steps", str(kwargs["episode_steps"])])
+
+        # find unused port
+        while port_in_use(self.base_port):
+            self.base_port += 1
+
+        # create logger
+        self.log = Logger(f"{kwargs['condition'].replace('-', '_')}{kwargs['run_id']}-{mode}",
+                          log_dir=f"{kwargs['log_path']}/")
+
+        # create environment and connect it to logger
+        self.env = UnityEnvironment(self.executable_path, side_channels=[self.log], additional_args=args, base_port=self.base_port)
+        self.env = UnityToGymWrapper(self.env, uint8_visual=True)
+
+        # initialize the parent class (gym.Wrapper)
+        super().__init__(self.env)
+
+    # converts the (c, w, h) frame returned by mlagents v1.0.0 and Unity 2022.3 to (w, h, c)
+    # as expected by gym==0.21.0
+    # HACK: mode is not used, but is required by the gym.Wrapper class (might be unnecessary but keeping for now)
+    def render(self, mode="rgb_array"): # pylint: disable=unused-argument
+        """
+        Renders the current frame of the environment.
+
+        :param mode: The mode to render the frame in. Defaults to "rgb_array".
+        :type mode: str, optional
+        :return: The rendered frame of the environment.
+        :rtype: np.ndarray
+        """
+        return np.moveaxis(self.env.render(), [0, 1, 2], [2, 0, 1])
+
+    def step(self, action):
+        """
+        Takes a step in the environment with the given action.
+
+        :param action: The action to take in the environment.
+        :type action: Any
+
+        :return: A tuple containing the next state, reward, done flag, and info dictionary.
+        :rtype: tuple[np.ndarray, float, bool, dict]
+        """
+        next_state, reward, done, info = self.env.step(action)
+        return next_state, float(reward), done, info
+
+    def log(self, msg: str) -> None:
+        """
+        Logs a message to the environment.
+
+        :param msg: The message to log.
+        :type msg: str
+        """
+        self.log.log_str(msg)
+
+    def reset(self, seed: Optional[int] = None, **kwargs): # pylint: disable=unused-argument
+        # nothing to do if the wrapped env does not accept `seed`
+        """
+        Resets the environment with the given seed and arguments.
+
+        :param seed: The seed to use for the environment. Defaults to None.
+        :type seed: int, optional
+        :param kwargs: The arguments to pass to the environment.
+        :type kwargs: Any
+
+        :return: The initial state of the environment.
+        :rtype: np.ndarray
+        """
+        return self.env.reset(**kwargs)
+
+    def __repr__(self) -> str:
+        attrs = {k: v for k, v in vars(self).items() if k != "logger"}
+        return f"{self.__class__.__name__}({attrs!r})"
+
+    def __str__(self) -> str:
+        attrs = {k: v for k, v in vars(self).items() if k != "logger"}
+        return f"{self.__class__.__name__}({attrs!r})"
diff --git a/src/nett/environment/configs.py b/src/nett/environment/configs.py
new file mode 100644
index 00000000..deda72e4
--- /dev/null
+++ b/src/nett/environment/configs.py
@@ -0,0 +1,289 @@
+"""This module contains the NETT configurations for different experiments."""
+
+import sys
+import inspect
+from abc import ABC, abstractmethod
+from itertools import product
+
+# the naming is confusing since it is used for train or test too.
+class NETTConfig(ABC):
+    """Abstract base class for NETT configurations.
+
+    :param param_defaults: A dictionary of parameter defaults.
+    :type param_defaults: dict[str, str]
+    :param **params: Keyword arguments representing the configuration parameters.
+    :raises ValueError: If any parameter value is not a value or subset of the default values.
+    """
+
+    def __init__(self, param_defaults: dict[str, str], **params) -> None:
+        """Constructor method
+        """
+        self.param_defaults = param_defaults
+        self.params = self._validate_params(params)
+        self.conditions = self._create_conditions_from_params(self.params)
+
+    def _create_conditions_from_params(self, params: dict[str, str]) -> list[str]:
+        """
+        Creates conditions from the configuration parameters.
+
+        :param params: The configuration parameters.
+        :type params: dict[str, str]
+        :return: A list of conditions.
+        :rtype: list[str]
+        """
+        combination_params = list(product(*params.values()))
+        conditions = ["-".join(combination).lower() for combination in combination_params]
+        return conditions
+
+    def _normalize_params(self, params: dict[str, str | int | float]) -> dict[str, str]:
+        """
+        Normalizes the configuration parameters.
+
+        :param params: The configuration parameters.
+        :type params: dict[str, str | int | float]
+        :return: The normalized configuration parameters.
+        :rtype: dict[str, str]
+
+        """
+        params = {param: (value if isinstance(value, list) else [value]) for param, value in params.items()}
+        params = {param: [str(item) for item in value] for param, value in params.items()}
+        return params
+
+    def _validate_params(self, params: dict[str, str]):
+        """
+        Validates the configuration parameters.
+
+        :param params: The configuration parameters.
+        :type params: dict[str, str]
+        :return: The validated configuration parameters.
+        :rtype: dict[str, str]
+        :raises ValueError: If any parameter value is not a value or subset of the default values.
+        """
+        params = self._normalize_params(params)
+        for (values, default_values) in zip(params.values(), self.param_defaults.values()):
+            if not set(values) <= set(default_values):
+                raise ValueError(f"{values} should be a value or subset of {default_values}")
+        return params
+
+    @property
+    def defaults(self):
+        """
+        Get the default values of the configuration parameters.
+
+        :return: A dictionary of parameter defaults.
+        :rtype: dict[str, Any]
+        """
+        signature = inspect.signature(self.__init__)
+        return {param: value.default for param, value in signature.parameters.items()
+                if value.default is not inspect.Parameter.empty}
+
+    @property
+    @abstractmethod
+    def num_conditions(self):
+        """
+        Get the number of conditions for the configuration.
+
+        :return: The number of conditions.
+        :rtype: int
+        """
+        pass
+
+
+class IdentityAndView(NETTConfig):
+    """
+    NETT configuration for Identity and View.
+
+    :param object: The object(s) to be used. Defaults to ["object1", "object2"].
+    :type object: str or list[str]
+    :param rotation: The rotation(s) to be used. Defaults to ["horizontal", "vertical"].
+    :type rotation: str or list[str]
+    """
+
+    def __init__(self,
+                 object: str | list[str] = ["object1", "object2"],
+                 rotation: str | list[str] = ["horizontal", "vertical"]) -> None:
+        """Constructor method
+        """
+        super().__init__(param_defaults=self.defaults,
+                         object=object,
+                         rotation=rotation)
+
+    @property
+    def num_conditions(self):
+        """
+        Get the number of conditions for the configuration.
+
+        :return: The number of conditions.
+        :rtype: int
+        """
+        return 18
+
+
+class Binding(NETTConfig):
+    """
+    NETT configuration for Binding.
+
+    :param object: The object(s) to be used. Defaults to ["object1", "object2"].
+    :type object: str or list[str]
+    """
+
+    def __init__(self,
+                 object: str | list[str] = ["object1", "object2"]) -> None:
+        """Constructor method
+        """
+        super().__init__(param_defaults=self.defaults,
+                         object=object)
+
+    @property
+    def num_conditions(self):
+        """
+        Get the number of conditions for the configuration.
+
+        :return: The number of conditions.
+        :rtype: int
+        """
+        return 50
+
+
+class Parsing(NETTConfig):
+    """
+    NETT configuration for Parsing.
+
+    :param background: The background(s) to be used. Defaults to ["A", "B", "C"].
+    :type background: str or list[str]
+    :param object: The object(s) to be used. Defaults to ["ship", "fork"].
+    :type object: str or list[str]
+    """
+
+    def __init__(self,
+                 background: str | list[str] = ["A", "B", "C"],
+                 object: str | list[str] = ["ship", "fork"]) -> None:
+        """Constructor method
+        """
+        super().__init__(param_defaults=self.defaults,
+                         background=background,
+                         object=object)
+
+    @property
+    def num_conditions(self):
+        """
+        Get the number of conditions for the configuration.
+
+        :return: The number of conditions.
+        :rtype: int
+        """
+        return 56
+
+
+class Slowness(NETTConfig):
+    """
+    NETT configuration for Slowness.
+
+    :param experiment: The experiment(s) to be used. Defaults to [1, 2].
+    :type experiment: str or list[int]
+    :param object: The object(s) to be used. Defaults to ["obj1", "obj2"].
+    :type object: str or list[str]
+    :param speed: The speed(s) to be used. Defaults to ["slow", "med", "fast"].
+    :type speed: str or list[str]
+    """
+
+    def __init__(self,
+                 experiment: str | list[int] = [1, 2],
+                 object: str | list[str] = ["obj1", "obj2"],
+                 speed: str | list[str] = ["slow", "med", "fast"]) -> None:
+        """Constructor method
+        """
+        super().__init__(param_defaults=self.defaults,
+                         experiment=experiment,
+                         object=object,
+                         speed=speed)
+
+    @property
+    def num_conditions(self):
+        """
+        Get the number of conditions for the configuration.
+
+        :return: The number of conditions.
+        :rtype: int
+        """
+        if self.params["experiment"] == "1":
+            return 5
+        return 13
+
+
+class Smoothness(NETTConfig):
+    """
+    NETT configuration for Smoothness.
+
+    :param object: The object(s) to be used. Defaults to ["obj1"].
+    :type object: str or list[str]
+    :param temporal: The temporal condition(s) to be used. Defaults to ["norm", "scram"].
+    :type temporal: str or list[str]
+    """
+
+    def __init__(self,
+                 object: str | list[str] = ["obj1"],
+                 temporal: str | list[str] = ["norm", "scram"]) -> None:
+        """Constructor method
+        """
+        super().__init__(param_defaults=self.defaults,
+                         object=object,
+                         temporal=temporal)
+
+    @property
+    def num_conditions(self):
+        """
+        Get the number of conditions for the configuration.
+
+        :return: The number of conditions.
+        :rtype: int
+        """
+        return 5
+
+
+class OneShotViewInvariant(NETTConfig):
+    """
+    NETT configuration for One-Shot View Invariant.
+
+    :param object: The object(s) to be used. Defaults to ["fork", "ship"].
+    :type object: str or list[str]
+    :param range: The range(s) to be used. Defaults to ["360", "small", "1"].
+    :type range: str or list[str]
+    :param view: The view(s) to be used. Defaults to ["front", "side"].
+    :type view: str or list[str]
+    """
+
+    def __init__(self,
+                 object: str | list[str] = ["fork", "ship"],
+                 range: str | list[str] = ["360", "small", "1"],
+                 view: str | list[str] = ["front", "side"]) -> None:
+        """Constructor method
+        """
+        super().__init__(param_defaults=self.defaults,
+                         object=object,
+                         range=range,
+                         view=view)
+
+    @property
+    def num_conditions(self):
+        """
+        Get the number of conditions for the configuration.
+
+        :return: The number of conditions.
+        :rtype: int
+        """
+        return 50
+
+
+def list_configs() -> set[str]:
+    """
+    Lists all available NETT configurations.
+
+    :return: A set of configuration names.
+    :rtype: set[str]
+
+    """
+    is_class_member = lambda member: inspect.isclass(member) and member.__module__ == __name__
+    clsmembers = inspect.getmembers(sys.modules[__name__], is_class_member)
+    clsmembers = [clsmember[0] for clsmember in clsmembers if clsmember[0] != "NETTConfig"]
+    return clsmembers
diff --git a/src/nett/nett.py b/src/nett/nett.py
new file mode 100644
index 00000000..5631db14
--- /dev/null
+++ b/src/nett/nett.py
@@ -0,0 +1,446 @@
+"""
+This module contains the NETT class, which is the main class for training, testing and analyzing brains in environments.
+
+.. module:: nett
+   :synopsis: Main class for training, testing and analyzing brains in environments.
+
+"""
+
+import time
+import subprocess
+from pathlib import Path
+from typing import Any
+from copy import deepcopy
+from itertools import product
+from concurrent.futures import ProcessPoolExecutor, Future
+
+import pandas as pd
+from sb3_contrib import RecurrentPPO
+from pynvml import nvmlInit, nvmlDeviceGetCount, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
+
+from nett import Brain, Body, Environment
+from nett.utils.io import mute
+
+
+class NETT:
+    """
+    The NETT class is the main class for training, testing and analyzing brains in environments.
+
+    :param brain: The brain to be trained and tested.
+    :type brain: Brain
+    :param body: The body to be used for training and testing the brain.
+    :type body: Body
+    :param environment: The environment in which the brain is to be trained and tested.
+    :type environment: Environment
+    
+    :ivar output_dir: The directory where the run results will be stored.
+    :vartype output_dir: Path
+    :ivar mode: The mode in which the brains are to be trained and tested. It can be "train", "test", or "full".
+    :vartype mode: str
+    :ivar verbosity: The verbosity level of the run.
+    :vartype verbosity: int
+    :ivar num_brains: The number of brains to be trained and tested.
+    :vartype num_brains: int
+    :ivar train_eps: The number of episodes the brains are to be trained for.
+    :vartype train_eps: int
+    :ivar test_eps: The number of episodes the brains are to be tested for.
+    :vartype test_eps: int
+    :ivar description: A description of the run.
+    :vartype description: str
+    :ivar buffer: The buffer for memory allocation.
+    :vartype buffer: float
+    :ivar step_per_episode: The number of steps per episode.
+    :vartype step_per_episode: int
+    :ivar device_type: The type of device to be used for training and testing. It can be "cuda" or "cpu".
+    :vartype device_type: str
+    :ivar devices: The list of devices to be used for training and testing. If -1, all available devices will be used.
+    :vartype devices: list[int] | int
+    :ivar brain: The brain to be trained and tested.
+    :vartype brain: Brain
+    :ivar body: The body to be used for training and testing the brain.
+    :vartype body: Body
+    :ivar environment: The environment in which the brain is to be trained and tested.
+    :vartype environment: Environment
+    :ivar logger: The logger for the NETT class.
+    :vartype logger: Logger
+
+    Example:
+
+    >>> from nett import NETT
+    >>> # create a brain, body, and environment
+    >>> benchmarks = NETT(brain, body, environment)
+    """
+
+    def __init__(self, brain: Brain, body: Body, environment: Environment) -> None:
+        """
+        Initialize the NETT class.
+
+        :param brain: The brain to be trained and tested.
+        :type brain: Brain
+        :param body: The body to be used for training and testing the brain.
+        :type body: Body
+        :param environment: The environment in which the brain is to be trained and tested.
+        :type environment: Environment
+        """
+        from nett import logger
+        self.logger = logger.getChild(__class__.__name__)
+        self.brain = brain
+        self.body = body
+        self.environment = environment
+        # for NVIDIA memory management
+        nvmlInit()
+
+    def run(self,
+            output_dir: Path | str,
+            num_brains: int = 1,
+            mode: str = "full",
+            train_eps: int = 1000,
+            test_eps: int = 20,
+            device_type: str = "cuda",
+            devices: list[int] | int =  -1,
+            description: str = None,
+            buffer: float = 1.2,
+            step_per_episode: int = 200,
+            verbosity: int = 0) -> list[Future]: # pylint: disable=unused-argument
+        """
+        Run the training and testing of the brains in the environment.
+
+        :param output_dir: The directory where the run results will be stored.
+        :type output_dir: Path | str
+        :param num_brains: The number of brains to be trained and tested.
+        :type num_brains: int, optional
+        :param mode: The mode in which the brains are to be trained and tested. It can be "train", "test", or "full".
+        :type mode: str, optional
+        :param train_eps: The number of episodes the brains are to be trained for.
+        :type train_eps: int, optional
+        :param test_eps: The number of episodes the brains are to be tested for.
+        :type test_eps: int, optional
+        :param device_type: The type of device to be used for training and testing. It can be "cuda" or "cpu".
+        :type device_type: str, optional
+        :param devices: The list of devices to be used for training and testing. If -1, all available devices will be used.
+        :type devices: list[int] | int, optional
+        :param description: A description of the run.
+        :type description: str, optional
+        :param buffer: The buffer for memory allocation.
+        :type buffer: float, optional
+        :param step_per_episode: The number of steps per episode.
+        :type step_per_episode: int, optional
+        :param verbosity: The verbosity level of the run.
+        :type verbosity: int, optional
+        
+        :return: A list of futures representing the jobs that have been launched.
+        :rtype: list[Future]
+
+        Example:
+    
+        >>> job_sheet = benchmarks.run(output_dir="./test_run", num_brains=2, train_eps=100, test_eps=10) # benchmarks is an instance of NETT
+        """
+        # set up the output_dir (wherever the user specifies, REQUIRED, NO DEFAULT)
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.logger.info(f"Set up run directory at: {self.output_dir.resolve()}")
+
+        # TODO (v0.3) upgrade to toml format
+        # register run config
+        self.mode = mode
+        self.verbosity = verbosity
+        self.num_brains = num_brains
+        self.train_eps = train_eps
+        self.test_eps = test_eps
+        self.description = description
+        self.buffer = buffer
+        self.step_per_episode = step_per_episode
+        self.device_type = self._validate_device_type(device_type)
+        self.devices = self._validate_devices(devices)
+
+        # schedule jobs
+        jobs = self._schedule_jobs()
+        self.logger.info("Scheduled jobs")
+
+        # launch jobs
+        self.logger.info("Launching")
+        job_sheet = self.launch_jobs(jobs)
+
+        # return control back to the user after launching jobs, do not block
+        return job_sheet
+
+    def launch_jobs(self, jobs: list[dict]) -> list[Future]:
+        """
+        Launch the jobs in the job sheet.
+
+        :param jobs: The jobs to be launched.
+        :type jobs: list[dict]
+        
+        :return: A list of futures representing the jobs that have been launched.
+        :rtype: list[Future]
+        """
+        max_workers = 1 if len(jobs) == 1 else None
+        initializer = mute if not self.verbosity else None
+        executor = ProcessPoolExecutor(max_workers=max_workers, initializer=initializer)
+        job_sheet = []
+        for job in jobs:
+            job_future = executor.submit(self._execute_job, job)
+            job_sheet.append({"running": job_future, "specification": job})
+            time.sleep(1)
+        return job_sheet
+
+    def status(self, job_sheet: dict[Future, dict]) -> pd.DataFrame:
+        """
+        Get the status of the jobs in the job sheet.
+        
+        :param job_sheet: The job sheet returned by the .launch_jobs() method.
+        :type job_sheet: dict[Future, dict]
+            
+        :return: A dataframe containing the status of the jobs in the job sheet.
+        :rtype: pd.DataFrame
+
+        Example:
+
+        >>> status = benchmarks.status(job_sheet) # benchmarks is an instance of NETT, job_sheet is the job sheet returned by the .run() method
+        """
+        selected_columns = ["brain_id", "condition", "device"]
+        filtered_job_sheet = [self._filter_job_record(job_record, selected_columns) for job_record in job_sheet]
+        return pd.json_normalize(filtered_job_sheet)
+
+# TODO v0.3, make .analyze() a staticmethod so that it does not need a class instance to call
+    # TODO v0.3. add support for user specified output_dir
+    # Discussion v0.3 is print okay or should we have it log using nett's logger?
+    # Discussion v0.3 move this out of the class entirely? from nett import analyze, analyze(...)
+
+    @staticmethod
+    def analyze(run_dir: str | Path,
+                output_dir: str | Path | None = None,
+                ep_bucket: int = 100,
+                num_episodes: int = 1000,
+                bar_order: str | list[int] = "default",
+                color_bars: bool = True) -> None:
+        """
+        Analyze the results of a run. This method is a static method and does not require an instance of the NETT class to be called.
+        
+        :param run_dir: The directory where the run results are stored.
+        :type run_dir: str | Path
+        :param output_dir: The directory where the analysis results will be stored. If None, the analysis results will be stored in the run directory.
+        :type output_dir: str | Path | None, optional
+        :param ep_bucket: The number of episodes to be grouped together for analysis.
+        :type ep_bucket: int, optional
+        :param num_episodes: The number of episodes to be analyzed.
+        :type num_episodes: int, optional
+        :param bar_order: The order in which the bars are to be displayed in the analysis plots. Default is "default". Can be "default", "asc", "desc", or a list of bar numbers (e.g. [3,1,2,4]).
+        :type bar_order: str | list[int], optional
+        :param color_bars: Whether to color the bars in the analysis plots by condition. Default is True.
+        :type color_bars: bool, optional
+            
+        :return: None
+        :rtype: None
+
+        Example:
+
+        >>> benchmarks.analyze(run_dir="./test_run", output_dir="./results") # benchmarks is an instance of NETT
+        """
+        # TODO may need to clean up this file structure
+        # set paths
+        run_dir = Path(run_dir).resolve()
+        analysis_dir = Path(__file__).resolve().parent.joinpath("analysis")
+        if output_dir is None:
+            output_dir = run_dir.joinpath("results")
+        output_dir = Path(output_dir).resolve()
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        # translate bar_order for R to read
+        bar_order_str = str(bar_order).translate({ord(i): None for i in ' []'}) # remove spaces and brackets from bar_order
+
+        # merge
+        print("Running merge")
+        subprocess.run(["Rscript", str(analysis_dir.joinpath("NETT_merge_csvs.R")),
+                        "--logs-dir", str(run_dir),
+                        "--results-dir", str(output_dir),
+                        "--results-name", "analysis_data",
+                        "--csv-train", "train_results.csv",
+                        "--csv-test", "test_results.csv"], check=True)
+
+        # train
+        print("Running analysis for [train]")
+        subprocess.run(["Rscript", str(analysis_dir.joinpath("NETT_train_viz.R")),
+                        "--data-loc", str(output_dir.joinpath("analysis_data")),
+                        "--results-wd", str(output_dir),
+                        "--ep-bucket", str(ep_bucket),
+                        "--num-episodes", str(num_episodes)], check=True)
+
+        # test
+        print("Running analysis for [test]")
+        subprocess.run(["Rscript", str(analysis_dir.joinpath("NETT_test_viz.R")),
+                        "--data-loc", str(output_dir.joinpath("analysis_data")),
+                        "--results-wd", str(output_dir),
+                        "--bar-order", bar_order_str,
+                        "--color-bars", str(color_bars),
+                        "--chick-file", str(analysis_dir.joinpath("ChickData", "ChickData_Parsing.csv"))], check=True)
+
+        print(f"Analysis complete. See results at {output_dir}")
+
+    def _schedule_jobs(self):
+        # get the free memory status for each device in list
+        free_device_memory = {device: memory_status["free"] for device, memory_status in self._get_memory_status().items()}
+
+        # estimate memory for a single job
+        job_memory = self._estimate_job_memory(free_device_memory)
+
+        # create jobs
+        # create list of all brain-environment combinations
+        # TODO (v0.5) replace environment.config.conditions and task_list with np.arrays to improve performance
+        task_list = list(product(self.environment.config.conditions, list(range(1, self.num_brains + 1))))
+
+        # assign devices based on memory to a condition and brain combination
+        # TODO (v0.5) what's better than a loop / brute force here?
+        jobs = []
+
+        free_devices = list(free_device_memory.keys()) # list of device numbers of free devices
+        dev_num = 0 # current device number
+
+        for (condition, brain_id) in task_list:
+            # find a device that has enough memory
+            # FIXME self.buffer does not seem like a reliable metric for estimating memory consumption (by multiplying 1.2 to job_memory)
+            while free_device_memory[free_devices[dev_num]] < job_memory * self.buffer:
+                dev_num += 1
+                # TODO (v0.3) allow partial execution
+                if dev_num >= len(free_device_memory):
+                    raise RuntimeError("Insufficient GPU Memory, could not create jobs for all tasks")
+            # allocate memory
+            free_device_memory[free_devices[dev_num]] -= job_memory*self.buffer
+            # create job
+            job = self._create_job(brain_id=brain_id, condition=condition, device=free_devices[dev_num])
+            jobs.append(job)
+
+        return jobs
+
+    def _create_job(self, device: int, condition: str, brain_id: int) -> dict:
+        # creates brain, env copies for a job
+        brain_copy = deepcopy(self.brain)
+
+        # configure paths for the job
+        paths = self._configure_job_paths(condition=condition, brain_id=brain_id)
+
+        # create job
+        return {"brain": brain_copy,
+                "environment": self.environment,
+                "body": self.body,
+                "device": device,
+                "condition": condition,
+                "brain_id": brain_id,
+                "paths": paths}
+
+    def _configure_job_paths(self, condition: str, brain_id: int) -> dict:
+        subdirs = ["model", "checkpoints", "plots", "logs", "env_recs", "env_logs"]
+        job_dir = Path.joinpath(self.output_dir, condition, f"brain_{brain_id}")
+        paths = {subdir: Path.joinpath(job_dir, subdir) for subdir in subdirs}
+        return paths
+
+    def _execute_job(self, job: dict[str, Any]) -> Future:
+        # for train
+        if self.mode not in ["train", "test", "full"]:
+            raise ValueError(f"Unknown mode type {self.mode}, should be one of ['train', 'test', 'full']")
+
+        # common environment kwargs
+        kwargs = {"rewarded": bool(self.brain.reward),
+                  "rec_path": str(job["paths"]["env_recs"]),
+                  "log_path": str(job["paths"]["env_logs"]),
+                  "condition": str(job["condition"]),
+                  "run_id": str(job["brain_id"]),
+                  "episode_steps": self.step_per_episode}
+
+        # for train
+        if self.mode in ["train", "full"]:
+            # initialize environment with necessary arguments
+            train_environment = deepcopy(job["environment"])
+            train_environment.initialize(mode="train", **kwargs)
+            # apply wrappers (body)
+            train_environment = job["body"](train_environment)
+            # train
+            job["brain"].train(env=train_environment,
+                               iterations=self.step_per_episode * self.train_eps,
+                               device_type=self.device_type,
+                               device=job["device"],
+                               paths=job["paths"])
+            train_environment.close()
+
+        # for test
+        if self.mode in ["test", "full"]:
+            # initialize environment with necessary arguments
+            test_environment = deepcopy(job["environment"])
+            test_environment.initialize(mode="test", **kwargs)
+            # apply wrappers (body)
+            test_environment = job["body"](test_environment)
+            # test
+            # readability over symmetry, sadly :(
+            if issubclass(job["brain"].algorithm, RecurrentPPO):
+                iterations = self.test_eps * test_environment.config.num_conditions
+            else:
+                iterations = self.step_per_episode * self.test_eps * job["environment"].config.num_conditions
+            job["brain"].test(env=test_environment,
+                              iterations=iterations,
+                              model_path=f"{job['paths']['model'].joinpath('latest_model.zip')}")
+            test_environment.close()
+
+        return f"Job Completed Successfully for Brain #{job['brain_id']} with Condition: {job['condition']}"
+
+    def _get_memory_status(self) -> dict[int, dict[str, int]]:
+        unpack = lambda memory_status: {"free": memory_status.free, "used": memory_status.used, "total": memory_status.total}
+        memory_status = {device_id : unpack(nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(device_id)))
+                         for device_id in self.devices}
+        return memory_status
+
+    # pylint: disable-next=unused-argument
+    def _estimate_job_memory(self, device_memory_status: dict) -> int: # pylint: disable=unused-argument
+        # TODO (v0.3) add a dummy job to gauge memory consumption
+
+        # # get device with the maxmium memory available
+        # max_memory_device = max(device_memory_status,
+        #                         key=lambda device: device_memory_status[device].free)
+        # # send a dummy job to gauge memory consumption
+        # dummy_job = self._create_job(device=max_memory_device,
+        #                              mode=random.choice(self.environment.config.modes),
+        #                              brain_id=-1)
+
+        # # send the job
+        # self.logger.info("Estimating memory by executing a single job")
+
+        # return a hurestic value for now (4GiB per job)
+        # multiply to return in bytes
+        memory_allocated = 4 * (1024 * 1024 * 1024)
+        return memory_allocated
+
+    def _filter_job_record(self, job_record: dict, selected_columns: list[str]) -> dict:
+        filter_specification = lambda specification: {key: value for key, value in specification.items() if key in selected_columns}
+        return {key: (filter_specification(value) if isinstance(value, dict) else value.running()) for key, value in job_record.items()}
+
+    def _validate_device_type(self, device_type: str):
+        # TODO (v0.4) add automatic type checking usimg pydantic or similar
+        if device_type in ["cuda", "cpu"]:
+            pass
+        else:
+            raise ValueError("Should be one of ['gpu', 'cpu']")
+        return device_type
+
+    def _validate_devices(self, devices: list[int] | int):
+        if self.device_type == "cpu" and isinstance(devices, list):
+            raise ValueError("Custom device lists not supported for 'cpu' device")
+
+        available_devices = list(range(nvmlDeviceGetCount()))
+        if isinstance(devices, list) and len(devices) > nvmlDeviceGetCount():
+            raise ValueError("Custom device lists not supported for 'cpu' device")
+
+        if devices == -1:
+            devices = available_devices
+            self.logger.info(f"Devices that will be used: {devices}")
+        elif isinstance(devices, list):
+            for device in devices:
+                if not isinstance(device, int) or device not in available_devices:
+                    raise ValueError(f"Device [{device}] is invalid. Available devices are: {available_devices}")
+            self.logger.info(f"Devices that will be used: {devices}")
+        else:
+            pass
+        return devices
+
+    def summary(self): # TODO: only raises a NotImplementedError for now
+        '''Generate a toml file and save it to the run directory.'''
+        raise NotImplementedError
+    
\ No newline at end of file
diff --git a/src/netts/utils/validation.py b/src/nett/utils/__init__.py
similarity index 100%
rename from src/netts/utils/validation.py
rename to src/nett/utils/__init__.py
diff --git a/src/netts/utils/callbacks.py b/src/nett/utils/callbacks.py
similarity index 85%
rename from src/netts/utils/callbacks.py
rename to src/nett/utils/callbacks.py
index fe5d3e59..6012cfb5 100644
--- a/src/netts/utils/callbacks.py
+++ b/src/nett/utils/callbacks.py
@@ -1,12 +1,19 @@
-import numpy as np
+"""
+Callbacks for training the agents.
 
+Classes:
+    HParamCallback(BaseCallback)
+    SupervisedSaveBestModelCallback(BaseCallback)
+"""
 from pathlib import Path
+import numpy as np
 from stable_baselines3.common.results_plotter import load_results, ts2xy
 from stable_baselines3.common.callbacks import BaseCallback
 from stable_baselines3.common.logger import HParam
-from netts.utils.train import compute_train_performance
 
-# TO DO (v0.3): refactor needed, especially logging
+from nett.utils.train import compute_train_performance
+
+# TODO (v0.3): refactor needed, especially logging
 class HParamCallback(BaseCallback):
     """
     Saves the hyperparameters and metrics at the start of the training, and logs them to TensorBoard.
@@ -33,8 +40,10 @@ def _on_training_start(self) -> None:
 
     def _on_step(self) -> bool:
         return True
-    
+
 class SupervisedSaveBestModelCallback(BaseCallback):
+    """
+    Callback to save the best model based on the mean performance of the last 100 episodes."""
     def __init__(self, summary_freq: int, save_dir: Path, env_log_path: str) -> None:
         super().__init__(verbose= 1)
         self.summary_freq = summary_freq
@@ -42,17 +51,17 @@ def __init__(self, summary_freq: int, save_dir: Path, env_log_path: str) -> None
         self.env_log_path = env_log_path
         self.best_mean_performance = -np.inf
         self.best_mean_reward = -np.inf
-        
+
     def _on_step(self) -> None:
         if self.n_calls % self.summary_freq == 0:
             # Retrieve training reward
             x, y = compute_train_performance(self.env_log_path)
             if len(x) > 0:
-                
+
                 # mean performance for last 100 episodes
                 mean_performance  = y[-1]
-                
-                x, y = ts2xy(load_results(self.env_log_path), 'timesteps')
+
+                x, y = ts2xy(load_results(self.env_log_path), "timesteps")
                 if len(x) > 0:
                     # mean reward for last 100 episodes
                     mean_reward  = np.mean(y[-100:])
@@ -62,7 +71,7 @@ def _on_step(self) -> None:
                             - Last mean reward per episode: {mean_reward:.2f}")
                     if mean_reward > self.best_mean_reward:
                         self.best_mean_reward = mean_reward
-            
+
                 if self.verbose > 0:
                     print(f"Best mean performance: {self.best_mean_performance:.2f}\
                         - Last mean performance per episode: {mean_performance:.2f}")
@@ -71,5 +80,5 @@ def _on_step(self) -> None:
                     if self.verbose > 0:
                         save_path = f"{self.save_dir.joinpath('best_model.zip')}"
                         print(f"Saving the best model to: {save_path}")
-                    self.model.save(save_path)    
-        return True
\ No newline at end of file
+                    self.model.save(save_path)
+        return True
diff --git a/src/netts/utils/environment.py b/src/nett/utils/environment.py
similarity index 57%
rename from src/netts/utils/environment.py
rename to src/nett/utils/environment.py
index be8988ad..accb7f5f 100644
--- a/src/netts/utils/environment.py
+++ b/src/nett/utils/environment.py
@@ -1,14 +1,25 @@
+"""
+environment.py
+
+Classes:
+    Logger(SideChannel)
+
+Functions:
+    port_in_use(port: int) -> bool
+"""
+
+import uuid
+import os
+import socket
 from mlagents_envs.side_channel.side_channel import (
     SideChannel,
     IncomingMessage,
     OutgoingMessage,
 )
-import uuid 
-import os 
-import socket
 
 
 def port_in_use(port):
+    """This function checks if a port is in use. It returns True if the port is in use and False if it is not."""
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     try:
         sock.bind(("localhost", port))
@@ -18,33 +29,42 @@ def port_in_use(port):
 
 # Create the StringLogChannel class. This is how logging info is communicated between python and unity
 class Logger(SideChannel):
+    """
+    This class is used to log information from the environment to a file. It inherits from the SideChannel class
+    
+    Methods:
+        on_message_received(msg: IncomingMessage) -> None
+        send_string(data: str) -> None
+        log_str(msg: str) -> None
+        __del__() -> None
+    """
     def __init__(self, log_title, log_dir="./EnvLogs/") -> None:
         super().__init__(uuid.UUID("621f0a70-4f87-11ea-a6bf-784f4387d1f7"))
         if not os.path.exists(log_dir):
             os.makedirs(log_dir)
         self.log_dir = log_dir
         f_name = os.path.join(log_dir, f"{log_title}.csv")
-        self.f = open(f_name, 'w')
+        self.f = open(f_name, "w")
 
-    #Method from Sidechannel interface. This method gets a message from unity
+    #Method from Sidechannel interface.
     def on_message_received(self, msg: IncomingMessage) -> None:
+        """This method is called when a message is received from unity."""
         self.f.write(msg.read_string()) #Write message to log file
         self.f.write("\n") #add new line character
 
-    #Method from Sidechannel interface. This method send a message to unity.
-    #This is here because it is required and I currently don't use it.
+    #This is here because it is required and I currently don"t use it.
     def send_string(self, data: str) -> None:
+        """Method from Sidechannel interface. This method send a message to unity."""
         msg = OutgoingMessage()
         msg.write_string(data)
         # We call this method to queue the data we want to send
         super().queue_message_to_send(msg)
-    
-    #This method writes a custom string to the log file
+
     def log_str(self, msg: str) -> None:
+        """This method writes a custom string to the log file"""
         self.f.write(msg)
         self.f.write("\n")
 
-    #This is called when the environment is shut down
     def __del__(self):
+        """This is called when the environment is shut down"""
         self.f.close()
-
diff --git a/src/nett/utils/io.py b/src/nett/utils/io.py
new file mode 100644
index 00000000..13808747
--- /dev/null
+++ b/src/nett/utils/io.py
@@ -0,0 +1,38 @@
+"""
+io
+
+This module contains utility functions for input/output operations
+
+Functions:
+    write_to_file: Write a dictionary to a file
+    mute: Mute the standard output and standard error
+"""
+import os
+import sys
+import json
+
+def write_to_file(file_path, d):
+    """
+    write_to_file
+    
+    Write a dictionary to a file
+    
+    Args:
+        file_path (str): The path to the file
+        d (dict): The dictionary to write
+        
+    Returns:
+        bool: True if the file was written, False otherwise
+    """
+    with open(file_path, "w") as file:
+        file.write(json.dumps(d))
+    return True
+
+def mute():
+    """
+    mute
+
+    Mute the standard output and standard error
+    """
+    sys.stdout = open(os.devnull, "w")
+    sys.stderr = open(os.devnull, "w") # TODO Should we be suppressing error messages?
diff --git a/src/netts/utils/logger.py b/src/nett/utils/logger.py
similarity index 52%
rename from src/netts/utils/logger.py
rename to src/nett/utils/logger.py
index 9726014e..e67d8145 100644
--- a/src/netts/utils/logger.py
+++ b/src/nett/utils/logger.py
@@ -1,42 +1,56 @@
+"""
+Logger
+
+Classes:
+    Logger
+"""
+
 #Packages for making the environment
+import uuid #needed for the communicator
+import os #Files and directories
+
 from mlagents_envs.side_channel.side_channel import (
     SideChannel,
     IncomingMessage,
     OutgoingMessage,
 )
-import uuid #needed for the communicator
-import os #Files and directories
-
 
-# Create the StringLogChannel class. This is how logging info is communicated between python and unity
+# Create the StringLogChannel class. 
 class Logger(SideChannel):
+    """
+    This class is used to log information from the environment to a file. This is how logging info is communicated between python and unity. It is a subclass of the SideChannel class from the mlagents_envs package.
+
+    Methods:
+        on_message_received: Method from Sidechannel interface. This method gets a message from unity and writes it to the log file.
+        send_string: Method from Sidechannel interface. This method send a message to unity.
+        log_str: This method is used to log a string to the file.
+    """
     def __init__(self, log_title, log_dir="./EnvLogs/") -> None:
-        super().__init__(uuid.UUID("621f0a70-4f87-11ea-a6bf-784f4387d1f7"))
+        super().__init__(uuid.UUID("621f0a70-4f87-11ea-a6bf-784f4387d1f7")) # TODO why this UUID?
         if not os.path.exists(log_dir):
             os.makedirs(log_dir)
         self.log_dir = log_dir
         f_name = os.path.join(log_dir, f"{log_title}.csv")
-        self.f = open(f_name, 'w')
+        self.f = open(f_name, "w")
 
-    #Method from Sidechannel interface. This method gets a message from unity
     def on_message_received(self, msg: IncomingMessage) -> None:
+        """Method from Sidechannel interface. This method gets a message from unity and writes it to the log file."""
         self.f.write(msg.read_string()) #Write message to log file
         self.f.write("\n") #add new line character
 
-    #Method from Sidechannel interface. This method send a message to unity.
     #This is here because it is required and I currently don't use it.
     def send_string(self, data: str) -> None:
+        """Method from Sidechannel interface. This method send a message to unity."""
         msg = OutgoingMessage()
         msg.write_string(data)
         # We call this method to queue the data we want to send
         super().queue_message_to_send(msg)
-    
-    #This method writes a custom string to the log file
+
     def log_str(self, msg: str) -> None:
+        """This method is used to log a string to the file."""
         self.f.write(msg)
         self.f.write("\n")
 
-    #This is called when the environment is shut down
     def __del__(self):
+        """This is called when the environment is shut down. It closes the log file."""
         self.f.close()
-
diff --git a/src/netts/utils/train.py b/src/nett/utils/train.py
similarity index 69%
rename from src/netts/utils/train.py
rename to src/nett/utils/train.py
index bc0d1831..8911ad77 100644
--- a/src/netts/utils/train.py
+++ b/src/nett/utils/train.py
@@ -1,32 +1,53 @@
+"""
+Train performance
+
+This module contains the functions to compute the train performance of the agent.
+
+Functions:
+    compute_train_performance: Compute Train performance
+    average_in_episode_three_region: Train performance
+    moving_average: Smooth values by doing a moving average
+"""
+
 import os
 import glob
 import numpy as np
 import pandas as pd
 
 def compute_train_performance(path):
+    """
+    Compute Train performance
+
+    Args:
+        path (Path or String?): path to training files # TODO check the type
+
+    Returns:
+        x (list): list of the episode numbers
+        y (np.array) : the moving averages of the success rate 
+    """
     x,y = [], []
     try:
         training_files = glob.glob(os.path.join(path, "*.csv"))
-        
+
         if len(training_files) == 0:
             raise Exception(f"Training file: {training_files} was not found in the {path}")
-        
-        
+
+
         for file_name in training_files:
-            
+
             log_df = pd.read_csv(file_name, skipinitialspace=True)
-            
-            percents,df,values = average_in_episode_three_region(log_df,"agent.x")
+
+            _, _, values = average_in_episode_three_region(log_df,"agent.x") # percents,df,
             y = moving_average(values, window=100)
-            x = list([i for i in range(0,len(y))])
-            
+            x = list(range(len(y)))
+
             break
-            
-        
+
+
         return x, y
     except Exception as ex:
         print(str(ex))
-        
+
     return x,y
 
 def average_in_episode_three_region(log,column='agent.x',transient=90):
@@ -35,7 +56,7 @@ def average_in_episode_three_region(log,column='agent.x',transient=90):
 
     Args:
         log (_type_): _description_
-        column (str, optional): _description_. Defaults to 'ChickAgent.x'.
+        column (str, optional): _description_. Defaults to 'agent.x'.
         transient (int, optional): _description_. Defaults to 90.
 
     Returns:
@@ -60,15 +81,13 @@ def average_in_episode_three_region(log,column='agent.x',transient=90):
             if np.isnan(percents[ep]):
                 percents[ep] = 0.5
 
-        rv = []
-        for key in percents:
-            rv.append(percents[key])
-        
+        rv = list(percents.values())
+
         return (percents,log,rv)
     except Exception as ex:
         print(str(ex))
         return (None, None)
-    
+
 def moving_average(values, window):
     """
     Smooth values by doing a moving average
@@ -77,4 +96,4 @@ def moving_average(values, window):
     :return: (numpy array)
     """
     weights = np.repeat(1.0, window) / window
-    return np.convolve(values, weights, 'valid')
\ No newline at end of file
+    return np.convolve(values, weights, 'valid')
diff --git a/src/netts/utils/visualization.py b/src/nett/utils/validation.py
similarity index 100%
rename from src/netts/utils/visualization.py
rename to src/nett/utils/validation.py
diff --git a/src/nett/utils/visualization.py b/src/nett/utils/visualization.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/netts/benchmark/benchmark.py b/src/netts/benchmark/benchmark.py
deleted file mode 100644
index 1fc25379..00000000
--- a/src/netts/benchmark/benchmark.py
+++ /dev/null
@@ -1,21 +0,0 @@
-import torch
-
-from netts import Brain, Body, Environment
-from netts.utils.gpu import getFirstAvailable
-
-class Benchmark:
-    def __init__(self, Brain, Body, Environment) -> None:
-        pass
-    
-    # TO DO (v0.3): Make memory dynamic, depending on the size of the model
-    # this is because different models require different GPU memory
-    def run(self, save_dir: str, device: list[int] | int | None = None):
-        self.save_dir = save_dir
-        self.device_num = getFirstAvailable(attempts=5, interval=5, maxMemory=0.5, verbose=True)
-        print(self.device_num)
-        torch.cuda.set_device(self.device_num[0])
-        assert torch.cuda.current_device() == self.device_num[0]
-
-    # TO DO (ideally before v1.0)
-    def publish(self):
-        raise NotImplementedError
\ No newline at end of file
diff --git a/src/netts/body/builder.py b/src/netts/body/builder.py
deleted file mode 100644
index 9b7e3fc4..00000000
--- a/src/netts/body/builder.py
+++ /dev/null
@@ -1,55 +0,0 @@
-from netts.body import types
-from netts.body import ascii_art
-
-from typing import Any
-from gym import Wrapper, Env
-
-# this will have the necessary wrappers before the observations interact with the brain, 
-# because it is the body that determines how the the observations will be processed.
-# specifically, in the case of the two-eyed agent, because the agent has two eyes, the observations are stereo 
-# and need to be processed differently before they make it to the brain. 
-# the body is the medium through which information travels from the environment to the brain. 
-# the brain is limited by what the body can percieve and no information is objective. 
-# NO INFORMATION IS OBJECTIVE (!!!!!!)
-class Body:
-    def __init__(self, type: str = "basic", 
-                 wrappers: list[Any] | None = None, 
-                 dvs: bool = False) -> None:
-        from netts import logger
-        self.logger = logger.getChild(__class__.__name__)
-        self.type = type
-        self.wrappers = self._validate_wrappers(wrappers)
-        self.dvs = self._validate_dvs(dvs)
-    
-    def _validate_agent_type(self, type: str) -> str:
-        if type not in types:
-            raise ValueError(f"agent_type must be one of {types}")
-        
-    def _validate_dvs(self, dvs: bool) -> bool:
-        if not isinstance(dvs, bool):
-            raise TypeError("dvs should be a boolean [True, False]")
-        return dvs
-    
-    def _validate_wrappers(self, wrappers: list[Any] | None) -> list[Any] | None:
-        if wrappers is not None:
-            for wrapper in wrappers:
-                if not isinstance(wrapper, Wrapper):
-                    raise ValueError("Wrappers must inherit from gym.Wrapper") 
-        return wrappers
-    
-    def __call__(self, env: Env) -> Env:
-        if self.wrappers:
-            for wrapper in self.wrappers:
-                env = wrapper(env)
-        return env
-
-    def __repr__(self) -> str:
-        attrs = {k: v for k, v in vars(self).items() if k != 'logger'}
-        return "{}({!r})".format(self.__class__.__name__, attrs)
-    
-    def __str__(self) -> str:
-        attrs = {k: v for k, v in vars(self).items() if k != 'logger'}
-        return "{}({!r})".format(self.__class__.__name__, attrs)
-    
-    def _register(self) -> None:
-        raise NotImplementedError
\ No newline at end of file
diff --git a/src/netts/brain/__init__.py b/src/netts/brain/__init__.py
deleted file mode 100644
index 4de7a68d..00000000
--- a/src/netts/brain/__init__.py
+++ /dev/null
@@ -1,62 +0,0 @@
-import ast
-import inspect
-import stable_baselines3
-import sb3_contrib
-
-from netts.brain import encoders
-from pathlib import Path
-from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
-
-# return all available enconders
-def list_encoders() -> set[str]:
-    encoder_dir = Path.joinpath(Path(__file__).resolve().parent, 'encoders')
-    encoders = [encoder.stem for encoder in list(encoder_dir.iterdir()) if "__" not in str(encoder)]
-    # set is faster to access than a list
-    return set(encoders)
-encoders_list = list_encoders()
-
-# return all available policy algorithms
-def list_algorithms() -> set[str]:
-    sb3_policy_algorithms = [algorithm for algorithm in dir(stable_baselines3) if algorithm[0].isupper()]
-    sb3_contrib_policy_algorithms = [algorithm for algorithm in dir(sb3_contrib) if algorithm[0].isupper()]
-    available_policy_algorithms = sb3_policy_algorithms + sb3_contrib_policy_algorithms
-    # set is faster to access than a list
-    return set(available_policy_algorithms)
-algorithms = list_algorithms()
-
-# TO DO (v0.3) return all available policy models programmatically
-def list_policies() -> set[str]:
-    return ['CnnPolicy', 'MlpPolicy', 'MultiInputPolicy', 'MultiInputLstmPolicy', 'CnnLstmPolicy']
-policies = list_policies()
-
-# return encoder string to encoder class mapping
-# TO DO (v0.3) optimized way to calculate and pass this dict around
-def get_encoder_dict():
-    encoder_dict = {}
-    encoder_dir = Path.joinpath(Path(__file__).resolve().parent, 'encoders')
-    # iterate through all files in the directory
-    for encoder_path in encoder_dir.iterdir():
-        if encoder_path.suffix == '.py' and "__" not in str(encoder_path):
-            module_name = encoder_path.stem 
-            # read the source
-            with open(encoder_path) as f:
-                source = f.read()
-            # parse it
-            module = ast.parse(source)
-            # get the first class definition
-            encoder_class = [node for node in ast.walk(module) if isinstance(node, ast.ClassDef)][0]
-            # add to the dictionary
-            encoder_dict[module_name] = encoder_class.name
-    return encoder_dict
-
-# TO DO (v0.3) return all available encoder classes programmatically
-# def get_encoder_dict():
-#     return {'cnnlstm': 'CNNLSTM', 
-#             'cotracker': 'CoTracker', 
-#             'dinov1': 'DinoV1',
-#             'dinov2': 'DinoV2',
-#             'vit': 'ViT',
-#             'resnet18': 'Resnet18CNN', 
-#             'sam': 'SegmentAnything', 
-#             'rnd': 'RND'}
-encoder_dict = get_encoder_dict()
\ No newline at end of file
diff --git a/src/netts/brain/builder.py b/src/netts/brain/builder.py
deleted file mode 100644
index facf1578..00000000
--- a/src/netts/brain/builder.py
+++ /dev/null
@@ -1,261 +0,0 @@
-import torch
-import gym
-import inspect
-import stable_baselines3
-import sb3_contrib
-import numpy as np
-import matplotlib.pyplot as plt
-
-from tqdm.auto import tqdm
-from typing import Any
-from pathlib import Path
-from netts.brain import algorithms, policies, encoders_list, encoder_dict
-from netts.brain import encoders
-from netts.utils.callbacks import SupervisedSaveBestModelCallback, HParamCallback
-from stable_baselines3.common.callbacks import CallbackList, CheckpointCallback
-from sb3_contrib import RecurrentPPO
-from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
-from stable_baselines3.common.policies import BasePolicy
-from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm
-from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
-from stable_baselines3.common.env_util import make_vec_env
-from stable_baselines3.common.env_checker import check_env
-from stable_baselines3.common.monitor import Monitor
-from stable_baselines3.common.logger import configure
-from stable_baselines3.common import results_plotter
-
-
-# TO DO (v0.2): Extend with support for custom policy models
-# TO DO (v0.2): should we move validation checks to utils under validations.py?
-class Brain:
-    def __init__(self, 
-                 encoder: Any | str = None,
-                 embedding_dim: int | None = None, 
-                 policy: Any | str | None = None,
-                 algorithm: Any | str | None = None,
-                 reward: Any | str = "supervised", 
-                 batch_size: int = 512, 
-                 buffer_size: int = 2048,
-                 train_encoder: bool | None = False,
-                 seed: int = 12
-                 ) -> None:
-        from netts import logger
-        self.logger = logger.getChild(__class__.__name__)
-        self.train_encoder = train_encoder
-        self.encoder = self._validate_encoder(encoder) if encoder else None
-        self.algorithm = self._validate_algorithm(algorithm) if algorithm else None
-        self.policy = self._validate_policy(policy) if policy else None
-        self.reward = self._validate_reward(reward) if reward else None
-        self.embedding_dim = embedding_dim
-        self.batch_size = batch_size
-        self.buffer_size = buffer_size
-        self.seed = seed
-
-    def train(self, env, iterations, device_type: str, device: int, paths: dict[str, Path]):
-        # validate environment
-        env = self._validate_env(env)
-
-        # initialize environment
-        log_path = paths['env_logs']
-        env = Monitor(env, str(log_path))
-        envs = make_vec_env(env_id=lambda : env, n_envs=1, seed=self.seed)
-
-        # build model
-        if self.encoder:
-            policy_kwargs = dict(features_extractor_class=self.encoder,
-                                 features_extractor_kwargs=dict(features_dim=inspect.signature(self.encoder).parameters['features_dim'].default))
-        else:
-            policy_kwargs = {}
-        self.model = self.algorithm(self.policy, 
-                                    envs, 
-                                    batch_size=self.batch_size,
-                                    n_steps=self.buffer_size,
-                                    verbose=1, 
-                                    policy_kwargs=policy_kwargs, 
-                                    device=torch.device(device_type, device))
-
-        # setup tensorboard logger and attach to model
-        tb_logger = configure(str(paths['logs']), ["stdout", "csv", "tensorboard"])
-        self.model.set_logger(tb_logger)
-        
-        # set encoder as eval only if train_encoder is not True
-        if not self.train_encoder:
-            self.model = self._set_encoder_as_eval(self.model)
-            self.logger.info(f'Encoder training is set to {str(self.train_encoder).upper()}')
-
-        # initialize callbacks
-        save_best_model_callback = SupervisedSaveBestModelCallback(summary_freq=2100,  
-                                                                   save_dir=paths['model'],
-                                                                   env_log_path=paths['env_logs'])
-        hparam_callback = HParamCallback()
-        checkpoint_callback = CheckpointCallback(save_freq=2100,
-                                                 save_path=paths["checkpoints"],
-                                                 name_prefix=self.algorithm.__name__,
-                                                 save_replay_buffer=True,
-                                                 save_vecnormalize=True)
-        callback_list = CallbackList([save_best_model_callback, hparam_callback, checkpoint_callback])
-
-        # train
-        self.logger.info(f"Total number of training steps: {iterations}")
-        self.model.learn(total_timesteps=iterations, 
-                         tb_log_name=self.algorithm.__name__, 
-                         progress_bar=True, 
-                         callback=[callback_list])
-        self.logger.info(f"Training Complete")
-        
-        # save
-        save_path = f"{paths['model'].joinpath('latest_model.zip')}"
-        self.save(save_path)
-        self.logger.info(f"Saved model at {save_path}")
-        # delete model to free up space
-        del self.model
-        # plot reward graph
-        self.plot_results(iterations=iterations, 
-                          model_log_dir=paths['env_logs'], 
-                          plots_dir=paths['plots'], 
-                          name="reward_graph")
-     
-    def test(self, env, iterations, model_path: str, record_prefix: str | None = None):
-        # load previously trained model from save_dir, if it exists
-        self.model = self.load(model_path)
-
-        # validate environment
-        env = self._validate_env(env)
-
-        # initialize environment
-        envs = make_vec_env(env_id=lambda : env, n_envs=1, seed=self.seed)
-        
-        # for when algorithm is RecurrentPPO
-        if issubclass(self.algorithm, RecurrentPPO):
-            self.logger.info(f'Testing with {self.algorithm.__name__}')
-            self.logger.info(f"Total number of episodes: {iterations}")
-            num_envs = 1
-            for episode in tqdm(range(iterations)):
-                obs = env.reset()
-                # cell and hidden state of the LSTM
-                done, lstm_states = False, None
-                # episode start signals are used to reset the lstm states
-                episode_starts = np.ones((num_envs,), dtype=bool)
-                episode_length = 0
-                while not done:
-                    action, lstm_states = self.model.predict(obs, state=lstm_states, episode_start=episode_starts, deterministic=True)
-                    obs, rewards, done, info = env.step(action)
-                    episode_starts = done
-                    episode_length += 1
-                    env.render(mode="rgb_array")
-
-        # for all other algorithms 
-        else: 
-            self.logger.info(f'Testing with {self.algorithm.__name__}')
-            self.logger.info(f"Total number of testing steps: {iterations}")
-            obs = envs.reset()
-            for step in tqdm(range(iterations)):
-                action, states = self.model.predict(obs, deterministic=True)
-                obs, reward, done, info = envs.step(action)
-                if done:
-                    env.reset()
-                env.render(mode="rgb_array")
-    
-    def save(self, path: str) -> None:
-        self.model.save(path)    
-
-    def load(self, model_path: str | Path):
-        return self.algorithm.load(model_path)
-    
-    def plot_results(self, iterations: int, model_log_dir: Path, plots_dir: Path, name: str) -> None:
-        results_plotter.plot_results([str(model_log_dir)], 
-                                     iterations, 
-                                     results_plotter.X_TIMESTEPS, 
-                                     name)
-        Path.mkdir(plots_dir)
-        plt.savefig(plots_dir.joinpath(f"{name}.png"))
-        plt.clf()
-
-    def _validate_encoder(self, encoder: Any | str) -> BaseFeaturesExtractor:
-        # for when encoder is a string
-        if isinstance(encoder, str):
-            if encoder not in encoder_dict.keys():
-                raise ValueError(f"If a string, should be one of: {encoder_dict.keys()}")
-            else:
-                encoder = getattr(encoders, encoder_dict[encoder])
-        
-        # for when encoder is a custom PyTorch encoder
-        if isinstance(encoder, BaseFeaturesExtractor):
-            # TO DO (v0.2) pass dummy torch.tensor on "meta" device to validate embedding dim
-            pass
-
-        if encoder and not hasattr(self, 'train_encoder'):
-            raise ValueError("encoder passed without setting train_encoder, should be one of: [True, False]")
-
-        return encoder
-        
-    def _validate_algorithm(self, algorithm: Any | str) -> OnPolicyAlgorithm | OffPolicyAlgorithm:
-        # for when policy is a string
-        if isinstance(algorithm, str):
-            if algorithm not in algorithms:
-                raise ValueError(f"If a string, should be one of: {algorithms}")
-            else:
-                # check for the passed policy in stable_baselines3 as well as sb3-contrib
-                # at this point in the code, it is guaranteed to be in either of the two
-                try:
-                    algorithm = getattr(stable_baselines3, algorithm)
-                except:
-                    algorithm = getattr(sb3_contrib, algorithm)
-       
-        # for when policy algorithm is custom 
-        elif isinstance(algorithm, OnPolicyAlgorithm) or isinstance(algorithm, OffPolicyAlgorithm):
-            # TO DO (v0.3) determine appropriate validation checks to be performed before passing
-            pass 
-        
-        else:
-            raise ValueError(f"Policy Algorithm should be either one of {algorithms} or a subclass of [{OnPolicyAlgorithm}, {OffPolicyAlgorithm}]")
-        
-        return algorithm
-    
-    def _validate_policy(self, policy: Any | str) -> str | BasePolicy:
-        # for when policy is a string
-        if isinstance(policy, str):
-            # support teseted for PPO and RecurrentPPO only
-            if policy not in policies:
-                raise ValueError(f"If a string, should be one of: {policies}")
-       
-        # for when policy is custom 
-        elif isinstance(policy, BasePolicy):
-            # TO DO (v0.3) determine appropriate validation checks to be performed before passing
-            pass 
-
-        else:
-            raise ValueError(f"Policy Model should be either one of {policies} or a subclass of [{BasePolicy}]")
-        
-        return policy
-    
-    def _validate_reward(self, reward: Any | str) -> Any | str:
-        # for when reward is a string
-        if isinstance(reward, str) and reward not in ['supervised', 'unsupervised']:
-            raise ValueError(f"If a string, should be one of: ['supervised', 'unsupervised']")
-        return reward
-    
-    # TO DO (v0.2) add typehinting for gym environments
-    def _validate_env(self, env) -> Any:
-        try:
-            check_env(env)
-        except Exception as ex:
-            raise ValueError(f"Failed training env check with {str(ex)}")
-        return env
-
-    def _set_encoder_as_eval(self, model):
-        model.policy.features_extractor.eval()
-        for param in model.policy.features_extractor.parameters():
-            param.requires_grad = False
-        return model
-
-    def __repr__(self) -> str:
-        attrs = {k: v for k, v in vars(self).items() if k != 'logger'}
-        return "{}({!r})".format(self.__class__.__name__, attrs)
-    
-    def __str__(self) -> str:
-        attrs = {k: v for k, v in vars(self).items() if k != 'logger'}
-        return "{}({!r})".format(self.__class__.__name__, attrs)
-    
-    
-    
\ No newline at end of file
diff --git a/src/netts/brain/encoders/dinov2.py b/src/netts/brain/encoders/dinov2.py
deleted file mode 100644
index 3fcc3d51..00000000
--- a/src/netts/brain/encoders/dinov2.py
+++ /dev/null
@@ -1,27 +0,0 @@
-import gym
-import torch
-
-from torchvision.transforms import Compose, Resize, CenterCrop, Normalize, InterpolationMode
-from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
-
-    
-class DinoV2(BaseFeaturesExtractor):
-    """
-    :param observation_space: (gym.Space)
-    :param features_dim: (int) Number of features extracted.
-        This corresponds to the number of unit for the last layer.
-    """
-    def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 384):
-        super(DinoV2, self).__init__(observation_space, features_dim)
-        self.n_input_channels = observation_space.shape[0]
-        self.transforms = Compose([Resize(size=256, 
-                                          interpolation=InterpolationMode.BICUBIC, 
-                                          max_size=None, 
-                                          antialias=True),
-                                   CenterCrop(size=(224, 224)),
-                                   Normalize(mean=torch.tensor([0.485, 0.456, 0.406]), 
-                                             std=torch.tensor([0.229, 0.224, 0.225]))])
-        self.model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14', pretrained=True)
-
-    def forward(self, observations: torch.Tensor) -> torch.Tensor:
-        return self.model(self.transforms(observations))
\ No newline at end of file
diff --git a/src/netts/brain/encoders/rnd.py b/src/netts/brain/encoders/rnd.py
deleted file mode 100644
index 808b5743..00000000
--- a/src/netts/brain/encoders/rnd.py
+++ /dev/null
@@ -1,61 +0,0 @@
-#!/usr/bin/env python3
-
-from typing import Dict, Tuple
-
-import torch as th
-from torch import nn
-from torch.nn import functional as F
-from torch.utils.data import DataLoader, TensorDataset
-
-
-class RND(nn.Module):
-    """Encoder for encoding observations.
-
-    Args:
-        obs_shape (Tuple): The data shape of observations.
-        action_dim (int): The dimension of actions.
-        latent_dim (int): The dimension of encoding vectors.
-
-    Returns:
-        Encoder instance.
-        
-    """
-    
-    def __init__(self, obs_shape: Tuple, action_dim: int, latent_dim: int) -> None:
-        super().__init__()
-        
-        ## visual
-        self.main = nn.Sequential(
-            nn.Conv2d(obs_shape[0], 32, kernel_size = 3, stride = 2, padding = 1),
-            nn.ELU(),
-            nn.Conv2d(32, 32, kernel_size = 3, stride = 2, padding = 1),
-            nn.ELU(),
-            nn.Conv2d(32, 32, kernel_size=3, stride=2, padding=1),
-            nn.ELU(),
-            nn.Conv2d(32, 32, kernel_size=3, stride=2, padding=1),
-            nn.ELU(),
-            nn.Flatten()
-        )
-        
-        with th.no_grad():
-            sample = th.ones(size = tuple(obs_shape))
-            n_flatten = self.main(sample.unsqueeze(0)).shape[1]
-        
-        self.linear = nn.Linear(n_flatten, latent_dim)
-        
-    def forward(self, obs: th.Tensor) -> th.Tensor:
-        """Encode the input tensors
-
-        Args:
-            obs (th.Tensor): Observations
-
-        Returns:
-            Encoding tensors
-        """
-        return self.linear(self.main(obs))
-    
-    
-    
-        
-     
-    
\ No newline at end of file
diff --git a/src/netts/brain/encoders/vit.py b/src/netts/brain/encoders/vit.py
deleted file mode 100644
index f15d18d2..00000000
--- a/src/netts/brain/encoders/vit.py
+++ /dev/null
@@ -1,27 +0,0 @@
-import gym
-import torch
-import timm
-
-from torchvision.transforms import Compose
-from torchvision.transforms import Resize, CenterCrop, Normalize, InterpolationMode
-from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
-
-class ViT(BaseFeaturesExtractor):
-    def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 384):
-        super(ViT, self).__init__(observation_space, features_dim)
-        self.n_input_channels = observation_space.shape[0]
-        self.transforms = Compose([Resize(size=248, 
-                                          interpolation=InterpolationMode.BICUBIC, 
-                                          max_size=None, 
-                                          antialias=True),
-                                   CenterCrop(size=(224, 224)),
-                                   Normalize(mean=torch.tensor([0.4850, 0.4560, 0.4060]), 
-                                             std=torch.tensor([0.2290, 0.2240, 0.2250]))])
-        
-        self.model = timm.create_model('vit_small_patch8_224.dino', 
-                                       in_chans=self.n_input_channels, 
-                                       num_classes=0, 
-                                       pretrained=False)
-                
-    def forward(self, observations: torch.Tensor) -> torch.Tensor:
-        return self.model(self.transforms(observations))
\ No newline at end of file
diff --git a/src/netts/environment/__init__.py b/src/netts/environment/__init__.py
deleted file mode 100644
index 9bfdacfc..00000000
--- a/src/netts/environment/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-# simplify imports
-from netts.environment.configs import list_configs
-from netts.environment.configs import IdentityAndView, Parsing, Slowness, Smoothness, Binding, OneShotViewInvariant
-
diff --git a/src/netts/environment/builder.py b/src/netts/environment/builder.py
deleted file mode 100644
index cef7e2d4..00000000
--- a/src/netts/environment/builder.py
+++ /dev/null
@@ -1,134 +0,0 @@
-from __future__ import annotations
-
-import os
-import numpy as np
-import subprocess
-
-from typing import Any
-from netts.environment.configs import NETTConfig, list_configs
-from netts.environment import configs
-from netts.utils.environment import Logger, port_in_use
-from gym import Wrapper
-from mlagents_envs.environment import UnityEnvironment
-from mlagents_envs.envs.unity_gym_env import UnityToGymWrapper
-
-class Environment(Wrapper):
-    def __init__(self, 
-                 config: str | NETTConfig, 
-                 executable_path: str, 
-                 display: int = 0,
-                 base_port: int = 5004,
-                 record_chamber: bool = False, 
-                 record_agent: bool = False, 
-                 recording_frames: int = 1000) -> None:
-        from netts import logger
-        self.logger = logger.getChild(__class__.__name__)
-        self.config = self._validate_config(config)
-        # TO DO (v0.4) what might be a way to check if it is a valid executable path? 
-        self.executable_path = executable_path
-        self.base_port = base_port
-        self.record_chamber = record_chamber
-        self.record_agent = record_agent
-        self.recording_frames = recording_frames
-        self.display = display
-
-        # set the correct permissions on the executable
-        self._set_executable_permission() 
-        # set the display for Unity environment
-        self._set_display()
-
-    def _validate_config(self, config: str | NETTConfig) -> NETTConfig:
-        # for when config is a str
-        if isinstance(config, str):
-            config_dict = {config_str.lower(): config_str for config_str in list_configs()}
-            if config not in config_dict.keys():
-                raise ValueError(f"Should be one of {config_dict.keys()}") 
-            else:
-                config = getattr(configs, config_dict[config])()
-
-        # for when config is a NETTConfig
-        elif isinstance(config, NETTConfig):
-            pass
-
-        else:
-            raise ValueError(f"Should either be one of {list(config_dict.keys())} or a subclass of NETTConfig")
-        
-        return config
-
-    def _set_executable_permission(self) -> None:
-        subprocess.run(['chmod', '-R', '755', self.executable_path])
-        self.logger.info("Executable permission is set")
-
-    def _set_display(self) -> None:
-        os.environ["DISPLAY"] = str(f":{self.display}")
-        self.logger.info("Display is set")
-
-
-    # copied from __init__() of chickai_env_wrapper.py (legacy)
-    # TO DO (v0.3) Critical refactor, don't like how this works, extremely error prone.  
-    # how can we build + constraint arguments better? something like an ArgumentParser sounds neat
-    # TO DO (v0.3) fix random_pos logic inside of Unity code
-    def initialize(self, **kwargs) -> Environment:
-        # from environment arguments
-        args = []
-        if self.recording_frames:
-            args.extend(["--recording-steps", str(self.recording_frames)])
-        if self.record_chamber:
-            args.extend(["--record-chamber", "true"])
-        if self.record_agent:
-            args.extend(["--record-agent", "true"])
-
-        # from runtime
-        if kwargs.get('mode', None):
-            args.extend(["--mode", kwargs['mode']])
-        if kwargs.get('rec_path', None):
-            args.extend(["--log-dir", f"{kwargs['rec_path']}/"])
-        # needs to fixed in Unity code where the default is always false
-        # if kwargs['mode'] != 'test':
-        #     args.extend(["--random-pos", "true"])
-        if kwargs.get("rewarded", False): 
-            args.extend(["--rewarded", "true"])
-        # TO DO: Discuss this with Manju, may be a MAJOR bug
-        self.step_per_episode = kwargs.get("episode_steps", 200)
-        if kwargs.get("episode_steps", False):
-            args.extend(["--episode-steps", str(kwargs["episode_steps"])])
-        
-        # find unused port 
-        while port_in_use(self.base_port):
-            self.base_port += 1
-
-        # create logger
-        self.log = Logger(f"{kwargs['run_id']}-{'train' if kwargs['mode'].split('-')[0] != 'test' else 'test'}", 
-                          log_dir=f"{kwargs['log_path']}/")
-        
-        # create environment and connect it to logger
-        self.env = UnityEnvironment(self.executable_path, side_channels=[self.log], additional_args=args, base_port=self.base_port)
-        self.env = UnityToGymWrapper(self.env, uint8_visual=True)
-        
-        # body?
-        # initialize the parent class (gym.Wrapper)
-        super().__init__(self.env)
-
-    # converts the (c, w, h) frame returned by mlagents v1.0.0 and Unity 2022.3 to (w, h, c) 
-    # as expected by gym==0.21.0
-    def render(self, mode="rgb_array"):
-        return np.moveaxis(self.env.render(), [0, 1, 2], [2, 0, 1])
-    
-    def step(self, action):
-        next_state, reward, done, info = self.env.step(action)
-        return next_state, float(reward), done, info
-    
-    def log(self, msg: str) -> None:
-        self.log.log_str(msg)
-
-    def reset(self, seed: Optional[int] = None, **kwargs):
-        # nothing to do if the wrapped env does not accept `seed`
-        return self.env.reset(**kwargs)
-
-    def __repr__(self) -> str:
-        attrs = {k: v for k, v in vars(self).items() if k != 'logger'}
-        return "{}({!r})".format(self.__class__.__name__, attrs)
-    
-    def __str__(self) -> str:
-        attrs = {k: v for k, v in vars(self).items() if k != 'logger'}
-        return "{}({!r})".format(self.__class__.__name__, attrs)
\ No newline at end of file
diff --git a/src/netts/environment/configs.py b/src/netts/environment/configs.py
deleted file mode 100644
index 3243c1ab..00000000
--- a/src/netts/environment/configs.py
+++ /dev/null
@@ -1,151 +0,0 @@
-import sys
-import inspect
-
-from abc import ABC, abstractmethod
-from itertools import product
-
-# TO DO: v0.3 change modes (picked up from legacy) to specs or conditions or similiar
-# the naming is confusing since it is used for train or test too. 
-class NETTConfig(ABC):
-    def __init__(self, param_defaults: dict[str, str], **params) -> None:
-        self.param_defaults = param_defaults
-        self.params = self._validate_params(params)
-        self.modes = self._create_modes_from_params(self.params)
-
-    def _create_modes_from_params(self, params: dict[str, str]) -> list[str]:
-        combination_params = list(product(*params.values()))
-        modes = ['-'.join(combination).lower() for combination in combination_params]
-        return modes
-    
-    def _normalize_params(self, params: dict[str, str | int | float]) -> dict[str, str]:
-        # add a list if a param has a singleton value (to make sure it works with itertools.product())
-        params = {param: (value if isinstance(value, list) else [value]) for param, value in params.items()}
-        # typecast values to str
-        params = {param: [str(item) for item in value] for param, value in params.items()}
-        return params
-    
-    def _validate_params(self, params: dict[str, str]):
-        params = self._normalize_params(params)
-        for (values, default_values) in zip(params.values(), self.param_defaults.values()):
-            if not set(values) <= set(default_values):
-                raise ValueError(f"{values} should be a value or subset of {default_values}")
-        return params
-
-    # don't quite like the round-robin happening here
-    # method of the super class is used to pass a value into the super class??
-    # move to utils?
-    @property
-    def defaults(self):
-        signature = inspect.signature(self.__init__)
-        return {param: value.default for param, value in signature.parameters.items() 
-                if value.default is not inspect.Parameter.empty}
-
-    @property
-    @abstractmethod
-    def num_conditions(self):
-        pass
-
-
-class IdentityAndView(NETTConfig):
-    """
-    NETT description goes here
-    """
-    def __init__(self,
-                 object: str | list[str] = ['object1', 'object2'], 
-                 rotation: str | list[str] = ['horizontal', 'vertical']) -> None:
-        super().__init__(param_defaults=self.defaults, 
-                         object=object, 
-                         rotation=rotation)
-
-    @property
-    def num_conditions(self):
-        return 18
-
-
-class Binding(NETTConfig):
-    """
-    NETT description goes here
-    """
-    def __init__(self,
-                 object: str | list[str] = ['object1', 'object2']) -> None:
-        super().__init__(param_defaults=self.defaults, 
-                         object=object)
-
-    @property
-    def num_conditions(self):
-        return 50
-
-class Parsing(NETTConfig):
-    """
-    NETT description goes here
-    """
-    def __init__(self,
-                 background: str | list[str] = ['A', 'B', 'C'], 
-                 object: str | list[str] = ['ship', 'fork']) -> None:
-        super().__init__(param_defaults=self.defaults, 
-                         background=background, 
-                         object=object)
-
-    @property
-    def num_conditions(self):
-        return 56
-
-class Slowness(NETTConfig):
-    """
-    NETT description goes here
-    """
-    def __init__(self,
-                 experiment: str | list[int] = [1, 2], 
-                 object: str | list[str] = ['obj1', 'obj2'],
-                 speed: str | list[str] = ['slow', 'med', 'fast']) -> None:
-        super().__init__(param_defaults=self.defaults, 
-                         experiment=experiment,
-                         object=object,
-                         speed=speed)
-
-    @property
-    def num_conditions(self):
-        if self.params['experiment'] == "1":
-            return 5
-        else:
-            return 13
-
-class Smoothness(NETTConfig):
-    """
-    NETT description goes here
-    """
-    def __init__(self,
-                 object: str | list[str] = ['obj1'],
-                 temporal: str | list[str] = ['norm', 'scram']) -> None:
-        super().__init__(param_defaults=self.defaults, 
-                         object=object,
-                         temporal=temporal)
-
-    @property
-    def num_conditions(self):
-        return 5
-
-class OneShotViewInvariant(NETTConfig):
-    """
-    NETT description goes here
-    """
-    def __init__(self,
-                 object: str | list[str] = ['fork', 'ship'],
-                 range: str | list[str] = ['360', 'small', '1'], 
-                 view: str | list[str] = ['front', 'side']) -> None:
-        super().__init__(param_defaults=self.defaults, 
-                         object=object,
-                         range=range,
-                         view=view)
-
-    @property
-    def num_conditions(self):
-        return 50
-    
-
-# get all available NETT configs
-def list_configs() -> set[str]:
-    is_class_member = lambda member: inspect.isclass(member) and member.__module__ == __name__
-    clsmembers = inspect.getmembers(sys.modules[__name__], is_class_member)
-    clsmembers = [clsmember[0] for clsmember in clsmembers if clsmember[0] != 'NETTConfig']
-    return clsmembers
\ No newline at end of file
diff --git a/src/netts/nett.py b/src/netts/nett.py
deleted file mode 100644
index bf060859..00000000
--- a/src/netts/nett.py
+++ /dev/null
@@ -1,236 +0,0 @@
-import time
-import pandas as pd
-
-from netts import Brain, Body, Environment
-from pathlib import Path
-from typing import Any
-from copy import deepcopy
-from dataclasses import dataclass
-from itertools import product
-from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, Future, as_completed
-from sb3_contrib import RecurrentPPO
-from pynvml import nvmlInit, nvmlDeviceGetCount, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
-
-
-class NETT:
-    def __init__(self, brain: Brain, body: Body, environment: Environment) -> None:
-        from netts import logger
-        self.logger = logger.getChild(__class__.__name__)
-        self.brain = brain
-        self.body = body
-        self.environment = environment     
-        # for NVIDIA memory management   
-        nvmlInit()
-
-    def run(self, 
-            dir: Path | str,
-            num_brains: int = 1, 
-            mode: str = "full", 
-            train_eps: int = 1000, 
-            test_eps: int = 20, 
-            device_type: str = "cuda",
-            devices: list[int] | int =  -1, 
-            description: str = None, 
-            buffer: float = 0.2,
-            step_per_episode: int = 200,
-            memory_per_brain: float = 0.5):
-        # set up the run_dir (wherever the user specifies, REQUIRED, NO DEFAULT)
-        self.run_dir = Path(dir)
-        self.run_dir.mkdir(parents=True, exist_ok=True)
-        self.logger.info(f"Set up run directory at: {self.run_dir.resolve()}")
-
-        # TO DO (v0.3) upgrade to toml format
-        # register run config
-        self.mode = mode
-        self.num_brains = num_brains
-        self.train_eps = train_eps
-        self.test_eps = test_eps
-        self.description = description
-        self.buffer = buffer
-        self.step_per_episode = step_per_episode
-        self.device_type = self._validate_device_type(device_type)
-        self.devices = self._validate_devices(devices)
-
-        # schedule jobs
-        jobs = self._schedule_jobs()
-        self.logger.info(f"Scheduled jobs")
-
-        # launch jobs
-        self.logger.info(f"Launching")
-        job_sheet = self.launch_jobs(jobs)
-        
-        # return control back to the user after launching jobs, do not block
-        return job_sheet
-    
-    def launch_jobs(self, jobs: list[dict]) -> list[Future]:
-        max_workers = 1 if len(jobs) == 1 else None
-        executor = ProcessPoolExecutor(max_workers=max_workers)
-        job_sheet = []
-        for job in jobs:
-            job_future = executor.submit(self._execute_job, job)
-            job_sheet.append({'running': job_future, 'specification': job})
-            time.sleep(1) # environment creation sometimes fails if no pause is given between consecutive job submissions
-        return job_sheet
-    
-    def status(self, job_sheet: dict[Future, dict]) -> pd.DataFrame:
-        selected_columns = ['brain_id', 'mode', 'device']
-        filtered_job_sheet = [self._filter_job_record(job_record, selected_columns) for job_record in job_sheet]
-        return pd.json_normalize(filtered_job_sheet)
-
-    def _schedule_jobs(self):
-        # get the free memory status for each device in list
-        free_device_memory = {device: memory_status['free'] for device, memory_status in self._get_memory_status().items()}
-
-        # estimate memory for a single job
-        job_memory = self._estimate_job_memory(free_device_memory)
-
-        # create jobs
-        task_list = list(product(self.environment.config.modes, list(range(1, self.num_brains + 1))))
-        # assign devices based on memory to a mode and brain combination
-        # TO DO (v0.5) what's better than a loop / brute force here? 
-        # TO DO (v0.3) replace explicit loop with .map() if does not sacrifice readability
-        jobs = []
-        # brute force iterate through devices
-        for device, free_memory in free_device_memory.items():
-            # iterate through modes and brain_ids
-            for (mode, brain_id) in task_list:
-                # check if memory is greater than job memory
-                if free_memory > job_memory * self.buffer:
-                    # alot device to task to make it a job which will be executed
-                    job = self._create_job(brain_id=brain_id, mode=mode, device=device)
-                    jobs.append(job)
-                    # remove task from list
-                    task_list.remove((mode, brain_id))
-                    # update free memory
-                    free_device_memory[device] = free_memory - job_memory
-                else:
-                    break
-        # ensure all tasks have been converted to jobs
-        # TO DO (v0.3) allow partial execution
-        if task_list:
-            raise RuntimeError("Insufficient GPU Memory, could not create jobs for all tasks")
-        return jobs
-    
-    def _create_job(self, device: int, mode: str, brain_id: int) -> dict:
-        # creates brain, env copies for a job
-        brain_copy = deepcopy(self.brain)
-
-        # configure paths for the job 
-        paths = self._configure_job_paths(mode=mode, brain_id=brain_id)
-
-        # create job
-        return {'brain': brain_copy, 
-                'environment': self.environment, 
-                'body': self.body, 
-                'device': device, 
-                'mode': mode,
-                'brain_id': brain_id,
-                'paths': paths}
-    
-    def _configure_job_paths(self, mode: str, brain_id: int) -> dict:
-        subdirs = ['model', 'checkpoints', 'plots', 'logs', 'env_recs', 'env_logs']
-        job_dir = Path.joinpath(self.run_dir, mode, f'brain_{brain_id}')
-        paths = {subdir: Path.joinpath(job_dir, subdir) for subdir in subdirs}
-        return paths
-
-    def _execute_job(self, job: dict[str, Any]) -> Future:
-        # common environment kwargs
-        kwargs = {'rewarded': True if self.brain.reward else False, 
-                  'rec_path': str(job['paths']['env_recs']), 
-                  'log_path': str(job['paths']['env_logs']), 
-                  'run_id': str(job['brain_id'])}
-        
-        # for train
-        if self.mode in ["train", "full"]:
-            # initialize environment with necessary arguments
-            train_environment = deepcopy(job['environment'])
-            train_environment.initialize(mode=f"train-{job['mode']}", **kwargs)
-            # apply wrappers (body)
-            train_environment = job['body'](train_environment)
-            # train
-            job['brain'].train(env=train_environment, 
-                               iterations=self.step_per_episode * self.train_eps, 
-                               device_type=self.device_type, 
-                               device=job['device'], 
-                               paths=job['paths'])
-
-        # for test
-        if self.mode in ["test", "full"]:
-            # initialize environment with necessary arguments
-            test_environment = deepcopy(job['environment'])
-            test_environment.initialize(mode=f"test-{job['mode']}", **kwargs)
-            # apply wrappers (body)
-            test_environment = job['body'](test_environment)
-            # test
-            # readability over symmetry, sadly :(
-            if issubclass(job['brain'].algorithm, RecurrentPPO):
-                iterations = self.test_eps * test_environment.config.num_conditions
-            else:
-                iterations = self.step_per_episode * self.test_eps * job['environment'].config.num_conditions
-            job['brain'].test(env=test_environment, 
-                              iterations=iterations, 
-                              model_path=f"{job['paths']['model'].joinpath('latest_model.zip')}")
-            
-        if self.mode not in ["train", "test", "full"]:
-            raise ValueError(f"Unknown mode type {self.mode}, should be one of ['train', 'test', 'full']")
-        
-        return "Job Successfully Concluded"
-    
-    def _get_memory_status(self) -> dict[int, dict[str, int]]:
-        unpack = lambda memory_status: {'free': memory_status.free, 'used': memory_status.used, 'total': memory_status.total}
-        memory_status = {device_id : unpack(nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(device_id))) 
-                         for device_id in self.devices}
-        return memory_status
-
-    def _estimate_job_memory(self, device_memory_status: dict) -> int:
-        # # get device with the maxmium memory available
-        # max_memory_device = max(device_memory_status, 
-        #                         key=lambda device: device_memory_status[device].free)
-        # # send a dummy job to gauge memory consumption
-        # dummy_job = self._create_job(device=max_memory_device, 
-        #                              mode=random.choice(self.environment.config.modes), 
-        #                              brain_id=-1)
-
-        # # send the job
-        # self.logger.info("Estimating memory by executing a single job")
-        
-        # return a hurestic value for now (4GiB per job)
-        # multiply to return in bytes
-        memory_allocated = 4 * (1024 * 1024 * 1024)
-        return memory_allocated
-    
-    def _filter_job_record(self, job_record: dict, selected_columns: list[str]) -> dict:
-        selected_columns = ['brain_id', 'mode', 'device']
-        filter_specification = lambda specification: {key: value for key, value in specification.items() if key in selected_columns}
-        return {key: (filter_specification(value) if isinstance(value, dict) else value.running()) for key, value in job_record.items()}
-    
-    def _validate_device_type(self, device_type: str):
-        # TO DO (v0.4) add automatic type checking usimg pydantic or similar
-        if device_type in ["cuda", "cpu"]:
-            pass
-        else:
-            raise ValueError("Should be one of ['gpu', 'cpu']")
-        return device_type
-
-    def _validate_devices(self, devices: list[int] | int):
-        if self.device_type == "cpu" and isinstance(devices, list):
-            raise ValueError("Custom device lists not supported for 'cpu' device")
-        else:
-            available_devices = list(range(nvmlDeviceGetCount()))
-            if isinstance(devices, list) and len(devices) > nvmlDeviceGetCount():
-                raise ValueError("Custom device lists not supported for 'cpu' device")
-            elif devices == -1:
-                devices = available_devices
-                self.logger.info(f"Devices that will be used: {devices}")
-            elif isinstance(devices, list):
-                for device in devices:
-                    if not isinstance(device, int) or device not in available_devices:
-                        raise ValueError(f"Device [{device}] is invalid. Available devices are: {available_devices}")
-                self.logger.info(f"Devices that will be used: {devices}")
-            else:
-                pass
-        return devices
-
-    # generate toml file and save to run directory
-    def summary():
-        raise NotImplementedError
\ No newline at end of file
diff --git a/src/netts/utils/gpu.py b/src/netts/utils/gpu.py
deleted file mode 100644
index b5bfc496..00000000
--- a/src/netts/utils/gpu.py
+++ /dev/null
@@ -1,312 +0,0 @@
-# GPUtil - GPU utilization
-#
-# A Python module for programmically getting the GPU utilization from NVIDA GPUs using nvidia-smi
-#
-# Author: Anders Krogh Mortensen (anderskm)
-# Date:   16 January 2017
-# Web:    https://github.com/anderskm/gputil
-#
-# LICENSE
-#
-# MIT License
-#
-# Copyright (c) 2017 anderskm
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-from subprocess import Popen, PIPE
-from distutils import spawn
-import os
-import math
-import random
-import time
-import sys
-import platform
-
-
-__version__ = '1.4.0'
-
-class GPU:
-    def __init__(self, ID, uuid, load, memoryTotal, memoryUsed, memoryFree, driver, gpu_name, serial, display_mode, display_active, temp_gpu):
-        self.id = ID
-        self.uuid = uuid
-        self.load = load
-        self.memoryUtil = float(memoryUsed)/float(memoryTotal)
-        self.memoryTotal = memoryTotal
-        self.memoryUsed = memoryUsed
-        self.memoryFree = memoryFree
-        self.driver = driver
-        self.name = gpu_name
-        self.serial = serial
-        self.display_mode = display_mode
-        self.display_active = display_active
-        self.temperature = temp_gpu
-
-def safeFloatCast(strNumber):
-    try:
-        number = float(strNumber)
-    except ValueError:
-        number = float('nan')
-    return number
-
-def getGPUs():
-    if platform.system() == "Windows":
-        # If the platform is Windows and nvidia-smi 
-        # could not be found from the environment path, 
-        # try to find it from system drive with default installation path
-        nvidia_smi = spawn.find_executable('nvidia-smi')
-        if nvidia_smi is None:
-            nvidia_smi = "%s\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe" % os.environ['systemdrive']
-    else:
-        nvidia_smi = "nvidia-smi"
-	
-    # Get ID, processing and memory utilization for all GPUs
-    try:
-        p = Popen([nvidia_smi,"--query-gpu=index,uuid,utilization.gpu,memory.total,memory.used,memory.free,driver_version,name,gpu_serial,display_active,display_mode,temperature.gpu", "--format=csv,noheader,nounits"], stdout=PIPE)
-        stdout, stderror = p.communicate()
-    except:
-        return []
-    output = stdout.decode('UTF-8')
-    # output = output[2:-1] # Remove b' and ' from string added by python
-    #print(output)
-    ## Parse output
-    # Split on line break
-    lines = output.split(os.linesep)
-    #print(lines)
-    numDevices = len(lines)-1
-    GPUs = []
-    for g in range(numDevices):
-        line = lines[g]
-        #print(line)
-        vals = line.split(', ')
-        #print(vals)
-        for i in range(12):
-            # print(vals[i])
-            if (i == 0):
-                deviceIds = int(vals[i])
-            elif (i == 1):
-                uuid = vals[i]
-            elif (i == 2):
-                gpuUtil = safeFloatCast(vals[i])/100
-            elif (i == 3):
-                memTotal = safeFloatCast(vals[i])
-            elif (i == 4):
-                memUsed = safeFloatCast(vals[i])
-            elif (i == 5):
-                memFree = safeFloatCast(vals[i])
-            elif (i == 6):
-                driver = vals[i]
-            elif (i == 7):
-                gpu_name = vals[i]
-            elif (i == 8):
-                serial = vals[i]
-            elif (i == 9):
-                display_active = vals[i]
-            elif (i == 10):
-                display_mode = vals[i]
-            elif (i == 11):
-                temp_gpu = safeFloatCast(vals[i]);
-        GPUs.append(GPU(deviceIds, uuid, gpuUtil, memTotal, memUsed, memFree, driver, gpu_name, serial, display_mode, display_active, temp_gpu))
-    return GPUs  # (deviceIds, gpuUtil, memUtil)
-
-
-def getAvailable(order = 'first', limit=1, maxLoad=0.5, maxMemory=0.5, memoryFree=0, includeNan=False, excludeID=[], excludeUUID=[]):
-    # order = first | last | random | load | memory
-    #    first --> select the GPU with the lowest ID (DEFAULT)
-    #    last --> select the GPU with the highest ID
-    #    random --> select a random available GPU
-    #    load --> select the GPU with the lowest load
-    #    memory --> select the GPU with the most memory available
-    # limit = 1 (DEFAULT), 2, ..., Inf
-    #     Limit sets the upper limit for the number of GPUs to return. E.g. if limit = 2, but only one is available, only one is returned.
-
-    # Get device IDs, load and memory usage
-    GPUs = getGPUs()
-
-    # Determine, which GPUs are available
-    GPUavailability = getAvailability(GPUs, maxLoad=maxLoad, maxMemory=maxMemory, memoryFree=memoryFree, includeNan=includeNan, excludeID=excludeID, excludeUUID=excludeUUID)
-    availAbleGPUindex = [idx for idx in range(0,len(GPUavailability)) if (GPUavailability[idx] == 1)]
-    # Discard unavailable GPUs
-    GPUs = [GPUs[g] for g in availAbleGPUindex]
-
-    # Sort available GPUs according to the order argument
-    if (order == 'first'):
-        GPUs.sort(key=lambda x: float('inf') if math.isnan(x.id) else x.id, reverse=False)
-    elif (order == 'last'):
-        GPUs.sort(key=lambda x: float('-inf') if math.isnan(x.id) else x.id, reverse=True)
-    elif (order == 'random'):
-        GPUs = [GPUs[g] for g in random.sample(range(0,len(GPUs)),len(GPUs))]
-    elif (order == 'load'):
-        GPUs.sort(key=lambda x: float('inf') if math.isnan(x.load) else x.load, reverse=False)
-    elif (order == 'memory'):
-        GPUs.sort(key=lambda x: float('inf') if math.isnan(x.memoryUtil) else x.memoryUtil, reverse=False)
-
-    # Extract the number of desired GPUs, but limited to the total number of available GPUs
-    GPUs = GPUs[0:min(limit, len(GPUs))]
-
-    # Extract the device IDs from the GPUs and return them
-    deviceIds = [gpu.id for gpu in GPUs]
-
-    return deviceIds
-
-#def getAvailability(GPUs, maxLoad = 0.5, maxMemory = 0.5, includeNan = False):
-#    # Determine, which GPUs are available
-#    GPUavailability = np.zeros(len(GPUs))
-#    for i in range(len(GPUs)):
-#        if (GPUs[i].load < maxLoad or (includeNan and np.isnan(GPUs[i].load))) and (GPUs[i].memoryUtil < maxMemory  or (includeNan and np.isnan(GPUs[i].memoryUtil))):
-#            GPUavailability[i] = 1
-
-def getAvailability(GPUs, maxLoad=0.5, maxMemory=0.5, memoryFree=0, includeNan=False, excludeID=[], excludeUUID=[]):
-    # Determine, which GPUs are available
-    GPUavailability = [1 if (gpu.memoryFree>=memoryFree) and (gpu.load < maxLoad or (includeNan and math.isnan(gpu.load))) and (gpu.memoryUtil < maxMemory  or (includeNan and math.isnan(gpu.memoryUtil))) and ((gpu.id not in excludeID) and (gpu.uuid not in excludeUUID)) else 0 for gpu in GPUs]
-    return GPUavailability
-
-def getFirstAvailable(order = 'first', maxLoad=0.5, maxMemory=0.5, attempts=1, interval=900, verbose=False, includeNan=False, excludeID=[], excludeUUID=[]):
-    #GPUs = getGPUs()
-    #firstAvailableGPU = np.NaN
-    #for i in range(len(GPUs)):
-    #    if (GPUs[i].load < maxLoad) & (GPUs[i].memory < maxMemory):
-    #        firstAvailableGPU = GPUs[i].id
-    #        break
-    #return firstAvailableGPU
-    for i in range(attempts):
-        if (verbose):
-            print('Attempting (' + str(i+1) + '/' + str(attempts) + ') to locate available GPU.')
-        # Get first available GPU
-        available = getAvailable(order=order, limit=1, maxLoad=maxLoad, maxMemory=maxMemory, includeNan=includeNan, excludeID=excludeID, excludeUUID=excludeUUID)
-        # If an available GPU was found, break for loop.
-        if (available):
-            if (verbose):
-                print('GPU ' + str(available) + ' located!')
-            break
-        # If this is not the last attempt, sleep for 'interval' seconds
-        if (i != attempts-1):
-            time.sleep(interval)
-    # Check if an GPU was found, or if the attempts simply ran out. Throw error, if no GPU was found
-    if (not(available)):
-        raise RuntimeError('Could not find an available GPU after ' + str(attempts) + ' attempts with ' + str(interval) + ' seconds interval.')
-
-    # Return found GPU
-    return available
-
-
-def showUtilization(all=False, attrList=None, useOldCode=False):
-    GPUs = getGPUs()
-    if (all):
-        if (useOldCode):
-            print(' ID | Name | Serial | UUID || GPU util. | Memory util. || Memory total | Memory used | Memory free || Display mode | Display active |')
-            print('------------------------------------------------------------------------------------------------------------------------------')
-            for gpu in GPUs:
-                print(' {0:2d} | {1:s}  | {2:s} | {3:s} || {4:3.0f}% | {5:3.0f}% || {6:.0f}MB | {7:.0f}MB | {8:.0f}MB || {9:s} | {10:s}'.format(gpu.id,gpu.name,gpu.serial,gpu.uuid,gpu.load*100,gpu.memoryUtil*100,gpu.memoryTotal,gpu.memoryUsed,gpu.memoryFree,gpu.display_mode,gpu.display_active))
-        else:
-            attrList = [[{'attr':'id','name':'ID'},
-                         {'attr':'name','name':'Name'},
-                         {'attr':'serial','name':'Serial'},
-                         {'attr':'uuid','name':'UUID'}],
-                        [{'attr':'temperature','name':'GPU temp.','suffix':'C','transform': lambda x: x,'precision':0},
-						 {'attr':'load','name':'GPU util.','suffix':'%','transform': lambda x: x*100,'precision':0},
-                         {'attr':'memoryUtil','name':'Memory util.','suffix':'%','transform': lambda x: x*100,'precision':0}],
-                        [{'attr':'memoryTotal','name':'Memory total','suffix':'MB','precision':0},
-                         {'attr':'memoryUsed','name':'Memory used','suffix':'MB','precision':0},
-                         {'attr':'memoryFree','name':'Memory free','suffix':'MB','precision':0}],
-                        [{'attr':'display_mode','name':'Display mode'},
-                         {'attr':'display_active','name':'Display active'}]]
-        
-    else:
-        if (useOldCode):
-            print(' ID  GPU  MEM')
-            print('--------------')
-            for gpu in GPUs:
-                print(' {0:2d} {1:3.0f}% {2:3.0f}%'.format(gpu.id, gpu.load*100, gpu.memoryUtil*100))
-        elif attrList is None:
-            # if `attrList` was not specified, use the default one
-            attrList = [[{'attr':'id','name':'ID'},
-                         {'attr':'load','name':'GPU','suffix':'%','transform': lambda x: x*100,'precision':0},
-                         {'attr':'memoryUtil','name':'MEM','suffix':'%','transform': lambda x: x*100,'precision':0}],
-                        ]
-        
-    if (not useOldCode):
-        if (attrList is not None):
-            headerString = ''
-            GPUstrings = ['']*len(GPUs)
-            for attrGroup in attrList:
-                #print(attrGroup)
-                for attrDict in attrGroup:
-                    headerString = headerString + '| ' + attrDict['name'] + ' '
-                    headerWidth = len(attrDict['name'])
-                    minWidth = len(attrDict['name'])
-                    
-                    attrPrecision = '.' + str(attrDict['precision']) if ('precision' in attrDict.keys()) else ''
-                    attrSuffix = str(attrDict['suffix']) if ('suffix' in attrDict.keys()) else ''
-                    attrTransform = attrDict['transform'] if ('transform' in attrDict.keys()) else lambda x : x
-                    for gpu in GPUs:
-                        attr = getattr(gpu,attrDict['attr'])
-                        
-                        attr = attrTransform(attr)
-                        
-                        if (isinstance(attr,float)):
-                            attrStr = ('{0:' + attrPrecision + 'f}').format(attr)
-                        elif (isinstance(attr,int)):
-                            attrStr = ('{0:d}').format(attr)
-                        elif (isinstance(attr,str)):
-                            attrStr = attr;
-                        elif  (sys.version_info[0] == 2):
-                            if (isinstance(attr,unicode)):
-                                attrStr = attr.encode('ascii','ignore')
-                        else:
-                            raise TypeError('Unhandled object type (' + str(type(attr)) + ') for attribute \'' + attrDict['name'] + '\'')
-                                            
-                        attrStr += attrSuffix
-                        
-                        minWidth = max(minWidth,len(attrStr))
-    
-                    headerString += ' '*max(0,minWidth-headerWidth)
-                    
-                    minWidthStr = str(minWidth - len(attrSuffix))
-                    
-                    for gpuIdx,gpu in enumerate(GPUs):
-                        attr = getattr(gpu,attrDict['attr'])
-                        
-                        attr = attrTransform(attr)
-                        
-                        if (isinstance(attr,float)):
-                            attrStr = ('{0:'+ minWidthStr + attrPrecision + 'f}').format(attr)
-                        elif (isinstance(attr,int)):
-                            attrStr = ('{0:' + minWidthStr + 'd}').format(attr)
-                        elif (isinstance(attr,str)):
-                            attrStr = ('{0:' + minWidthStr + 's}').format(attr);
-                        elif (sys.version_info[0] == 2):
-                            if (isinstance(attr,unicode)):
-                                attrStr = ('{0:' + minWidthStr + 's}').format(attr.encode('ascii','ignore'))
-                        else:
-                            raise TypeError('Unhandled object type (' + str(type(attr)) + ') for attribute \'' + attrDict['name'] + '\'')
-                                            
-                        attrStr += attrSuffix
-                        
-                        GPUstrings[gpuIdx] += '| ' + attrStr + ' '
-                                            
-                headerString = headerString + '|'
-                for gpuIdx,gpu in enumerate(GPUs):
-                    GPUstrings[gpuIdx] += '|'
-                    
-            headerSpacingString = '-' * len(headerString)
-            print(headerString)
-            print(headerSpacingString)
-            for GPUstring in GPUstrings:
-                print(GPUstring)
\ No newline at end of file
diff --git a/src/netts/utils/io.py b/src/netts/utils/io.py
deleted file mode 100644
index a70a6f38..00000000
--- a/src/netts/utils/io.py
+++ /dev/null
@@ -1,12 +0,0 @@
-import json 
-
-from pathlib import Path
-
-def write_to_file(file_path, d):
-    with open(file_path, 'w') as file:
-        file.write(json.dumps(d)) 
-    return True
-
-def configure_save_dir():
-    raise NotImplementedError
-


Test patch:
diff --git a/src/nett/analysis/NETT_test_viz.R b/src/nett/analysis/NETT_test_viz.R
new file mode 100644
index 00000000..52da1be1
--- /dev/null
+++ b/src/nett/analysis/NETT_test_viz.R
@@ -0,0 +1,229 @@
+#!/usr/bin/env Rscript
+
+# NETT_test_viz.R
+
+# Before running this script, you need to run merge_csvs to merge all of the agents'
+# output into a single, standardized format dataframe for training and test data
+
+# Variables --------------------------------------------------------------------
+
+# Read in the user-specified variables:
+library(argparse)
+parser <- ArgumentParser(description="An executable R script for the Newborn Embodied Turing Tests to analyze test trials")
+parser$add_argument("--data-loc", type="character", dest="data_loc",
+                    help="Full filename (inc working directory) of the merged R data",
+                    required=TRUE)
+parser$add_argument("--chick-file", type="character", dest="chick_file",
+                    help="Full filename (inc working directory) of the chick data CSV file",
+                    required=TRUE)
+parser$add_argument("--results-wd", type="character", dest="results_wd",
+                    help="Working directory to save the resulting visualizations",
+                    required=TRUE)
+parser$add_argument("--bar-order", type = "character", default = "default", dest = "bar_order",
+                    help="Order of bars. Use 'default', 'asc', 'desc', or specify indices separated by commas (e.g., '3,2,1,4')",
+                    required=FALSE)
+parser$add_argument("--color-bars", type = "character", dest="color_bars",
+                    help="Should the bars be colored by test condition?",
+                    required=TRUE)
+
+# Set script variables based on user input
+args <- parser$parse_args()
+data_loc <- args$data_loc; chick_file <- args$chick_file; results_wd <- args$results_wd; bar_order <- args$bar_order
+if( args$color_bars %in% c("t", "T", "true", "TRUE", "True")) {color_bars <- TRUE} else { color_bars <- FALSE}
+
+# Set Up -----------------------------------------------------------------------
+
+library(tidyverse)
+library(stringr)
+
+# Load the chick data
+chick_data <- read.csv(chick_file)
+
+# Load test data
+load(data_loc)
+rm(train_data)
+
+# Code each episode correct/incorrect
+test_data <- test_data %>%
+  mutate(correct_steps = if_else(correct.monitor == " left", left_steps, right_steps)) %>%
+  mutate(incorrect_steps = if_else(correct.monitor == " left", right_steps, left_steps)) %>%
+  mutate(percent_correct = correct_steps / (correct_steps + incorrect_steps))
+
+# Adjust bar order according to user input -------------------------------------
+
+# Create a variable to store the final order
+order <- NULL
+if (bar_order == "default" || bar_order == "asc" || bar_order == "desc"){
+  order <- bar_order
+}else {
+  order <- as.integer(strsplit(order_input, ",")[[1]])
+}
+
+# Conditionally reorder the dataframe based on user input
+if (!is.null(order)) {
+  if (order == "desc") {
+    test_data <- test_data %>%
+      arrange(desc(percent_correct)) %>%
+      mutate(test.cond = factor(test.cond, levels = unique(test.cond)))
+  } else if (order == "asc"){
+    test_data <- test_data %>%
+      arrange(percent_correct) %>%
+      mutate(test.cond = factor(test.cond, levels = unique(test.cond)))
+  } else if (order != "default") {
+    # Map numeric indices to factor levels
+    current_order <- levels(factor(test_data$test.cond))
+    new_order <- current_order[order]
+    test_data$test.cond <- factor(test_data$test.cond, levels = new_order)
+  }
+  # If order is "default", no need to change anything
+}
+
+
+# Plot aesthetic settings ------------------------------------------------------
+custom_palette <- c("#3F8CB7", "#FCEF88", "#5D5797", "#62AC6B", "#B74779", "#2C4E98","#CCCCE7", "#08625B", "#D15056")
+chickred <- "#AF264A"
+
+p <- ggplot() +
+  theme_classic() +
+  theme(axis.text.x = element_text(size = 6)) +
+  ylab("Percent Correct") +
+  xlab("Test Condition") +
+  scale_y_continuous(expand = c(0, 0), limits = c(0, 1), breaks=seq(0,1,.1), labels = scales::percent) +
+  geom_hline(yintercept = .5, linetype = 2) +
+  scale_fill_manual(values = custom_palette) +
+  scale_colour_manual(values = custom_palette) +
+  theme(axis.title = element_text(face="bold"),
+        axis.text.x = element_text(face="bold", size=7.5),
+        axis.text.y = element_text(face="bold", size=7.5))
+
+
+# Bar Chart Function -----------------------------------------------------------
+make_bar_charts <- function(data, dots, aes_y, error_min, error_max, img_name)
+{
+  p +
+
+    # Add chicken performance FIRST to sort the bars
+    geom_errorbar(data=chick_data, width = 0.7, colour = chickred,
+                  aes(x=test.cond, ymin=avg, ymax=avg)) +
+
+    # Model performance: bars
+    {if(color_bars)geom_col(data = data, width = 0.7, aes(x=test.cond, y = {{aes_y}}, fill = test.cond))}+
+    {if(!color_bars)geom_col(data = data, width = 0.7, aes(x=test.cond, y = {{aes_y}}), fill = "gray45")}+
+    # Model performance: error bars
+    geom_errorbar(data = data, width = 0.3,
+                  aes(x = test.cond, ymin = {{error_min}}, ymax = {{error_max}})) +
+    # Model performance: dots
+    {if(!is.null(dots))geom_jitter(data = dot_data, aes(x=test.cond, y = avgs), width = .3)}+
+    theme(legend.position="none") +
+
+    # Add chicken performance again so that it shows up on top
+    # Chick performance: lines (errorbar) with ribbons (crossbar)
+    geom_errorbar(data=chick_data, width = 0.7, colour = chickred,
+                aes(x=test.cond, ymin=avg, ymax=avg)) +
+    geom_crossbar(data=chick_data, width = 0.7, 
+                  linetype = 0, fill = chickred, alpha = 0.2,
+                  aes(x = test.cond, y = avg, 
+                      ymin = avg - avg_dev, ymax = avg + avg_dev)) 
+  
+  ggsave(img_name, width = 6, height = 6)
+}
+
+# Switch wd before we save the graphs
+setwd(results_wd)
+
+# Plot by agent ----------------------------------------------------------------
+## Leave rest data for agent-level graphs
+
+## Group data by test conditions
+by_test_cond <- test_data %>%
+  group_by(imprint.cond, agent, test.cond) %>%
+  summarise(avgs = mean(percent_correct, na.rm = TRUE), 
+            sd = sd(percent_correct, na.rm = TRUE), 
+            count = length(percent_correct),
+            tval = tryCatch({ (t.test(percent_correct, mu=0.5)$statistic)}, error = function(err){NA}),
+            df = tryCatch({(t.test(percent_correct, mu=0.5)$parameter)},error = function(err){NA}),
+            pval = tryCatch({(t.test(percent_correct, mu=0.5)$p.value)},error = function(err){NA}))%>%
+  mutate(se = sd / sqrt(count)) %>%
+  mutate(cohensd = (avgs - .5) / sd) %>%
+  mutate(imp_agent = paste(imprint.cond, agent, sep="_"))
+
+write.csv(by_test_cond, "stats_by_agent.csv")
+
+for (i in unique(by_test_cond$imp_agent))
+{
+  bar_data <- by_test_cond %>%
+    filter(imp_agent == i)
+
+  img_name <- paste0(i, "_test.png")
+
+  make_bar_charts(data = bar_data,
+                  dots = NULL,
+                  aes_y = avgs,
+                  error_min = avgs - se,
+                  error_max = avgs + se,
+                  img_name = img_name)
+}
+
+
+# Plot by imprinting condition -------------------------------------------------
+## Remove rest data once we start to group agents (for ease of presentation)
+
+by_imp_cond <- by_test_cond %>%
+  ungroup() %>%
+  group_by(imprint.cond, test.cond) %>%
+  summarise(avgs_by_imp = mean(avgs, na.rm = TRUE), 
+            sd = sd(avgs, na.rm = TRUE), 
+            count = length(avgs),
+            tval = tryCatch({ (t.test(avgs, mu=0.5)$statistic)}, error = function(err){NA}),
+            df = tryCatch({ (t.test(avgs, mu=0.5)$parameter)}, error = function(err){NA}),
+            pval = tryCatch({ (t.test(avgs, mu=0.5)$p.value)}, error = function(err){NA}))%>%
+  mutate(se = sd / sqrt(count)) %>%
+  mutate(cohensd = (avgs_by_imp - .5) / sd)
+
+write.csv(by_imp_cond, "stats_by_imp_cond.csv")
+
+for (i in unique(by_imp_cond$imprint.cond))
+{
+  bar_data <- by_imp_cond %>%
+    filter(imprint.cond == i) %>%
+    filter(test.cond != "Rest")
+  
+  dot_data <- by_test_cond %>%
+    filter(imprint.cond == i) %>%
+    filter(test.cond != "Rest")
+
+  img_name <- paste0(i, "_test.png")
+
+  make_bar_charts(data = bar_data,
+                  dots = dot_data,
+                  aes_y = avgs_by_imp,
+                  error_min = avgs_by_imp - se,
+                  error_max = avgs_by_imp + se,
+                  img_name = img_name)
+}
+
+
+# Plot across all imprinting conditions ----------------------------------------
+across_imp_cond <- by_test_cond %>%
+  ungroup() %>%
+  filter(test.cond != "Rest") %>%
+  group_by(test.cond) %>%
+  summarise(all_avgs = mean(avgs, na.rm = TRUE), 
+            sd = sd(avgs, na.rm = TRUE), 
+            count = length(avgs),
+            tval = tryCatch({ (t.test(avgs, mu=0.5)$statistic)}, error = function(err){NA}),
+            df =  tryCatch({ (t.test(avgs, mu=0.5)$parameter)}, error = function(err){NA}),
+            pval = tryCatch({ (t.test(avgs, mu=0.5)$p.value)}, error = function(err){NA}))%>%
+  mutate(se = sd / sqrt(count)) %>%
+  mutate(cohensd = (all_avgs - .5) / sd)
+
+write.csv(across_imp_cond, "stats_across_all_agents.csv")
+
+dot_data <- filter(by_test_cond, test.cond != "Rest")
+
+make_bar_charts(data = across_imp_cond,
+                dots = dot_data,
+                aes_y = all_avgs,
+                error_min = all_avgs - se,
+                error_max = all_avgs + se,
+                img_name = "all_imprinting_conds_test.png")