Environment setup + train loop (#5)

Mujoco environment setup (vectorized on GPU) + training loop + simulate script

docs/api/environment.md

@@ -0,0 +1,24 @@
+# Brittle star environment
+
+## Creation
+The environment package contains a factory class `BrittleStarEnvFactory` 
+that creates instances of the environment/morphologies/... It uses the
+configuration classes defined in `env_config.py` to create the instances.
+
+## Configuration
+The data classes in `env_config` have default values as stated in the tutorials.
+* MorphologyConfig: configuration for the morphology of the brittle star. Contains
+number of arms, number of segments per arm, and control mode.
+* ArenaConfig: configuration for the arena. Sets the size of the arena, whether to 
+set the ground floor to sand, attach a target and sizes of the walls.
+* EnvConfig: configuration for the environment. These set shared settings
+such as camera locations, simulation time and the task.
+
+## Backend and Task enums
+The Backend enum specifies either an MJC or MJX backend.
+* MJC: runs on CPU
+* MJX: uses jax on the gpu
+
+The Task enum specifies which task to use. 2 items are present:
+* DIRECTED_LOCOMOTION: move to a target location
+* LIGHT_ESCAPE: situation where the robot must move to a darker location

docs/api/train_simulate.md

@@ -0,0 +1,30 @@
+# Training and Simulation for Brittle Star Models
+
+## Training a model
+
+To train a model, you can use the `train.py` script. This script allows to pass some parameters to customize the training process:
+
+- `--out`: The output path where the trained model will be saved.
+- `--model_type`: The type of model to train (e.g., `random`, ...)
+- `--task`: The task to train on (e.g., `directed_locomotion`, ...)
+- `--seed`: The random seed for reproducibility.
+- `--epochs`: The number of epochs to train for.
+
+This will then train the specified model on the specified task for the given number of epochs and save the trained model to the specified output path.
+
+```bash
+python train.py --out artifacts/my_model --model-type random --task directed_locomotion --seed 0 --epochs 50 
+```
+
+## Simulating a model
+
+In order to simulate and view the behavior of a trained model, you can use the `simulate.py` script. This script allows you to specify the path to a trained model and will launch a simulation using that model. This script has the following parameters:
+
+- `--model`: The path to the trained model artifact to simulate.
+- `--model-type`: The type of model to simulate (e.g., `random`, ...)
+- `--task`: The task to simulate (e.g., `directed_locomotion`, ...)
+- `--seed`: The random seed for reproducibility.
+
+```bash
+python simulate.py --model artifacts/my_model --model-type random --task directed_locomotion --seed 0
+```