diff --git a/arduino/arduinocontrol/arduinocontrol.ino b/arduino/arduinocontrol/arduinocontrol.ino
index 2104156..f3ad4d3 100644
--- a/arduino/arduinocontrol/arduinocontrol.ino
+++ b/arduino/arduinocontrol/arduinocontrol.ino
@@ -1,274 +1,275 @@
-#include <Servo.h>
-
-// #define CALIBRATION_BOARD // uncomment this line if you are using the calibration pins
-// #define RPM_SENSOR // uncomment if you want to compile with the RPM_SENSOR
-// #define ENABLE_BRAKES // uncomment if you want to use the brakes
-
-// Servo
-#define SERVO_PIN 6
-#define SERVO_MIN 900
-#define SERVO_MAX 2100
-#define SERVO_NEUTRAL 1500
-#define SERVO_DEADBAND 2
-Servo servoSteering;
-
-// Motor
-#define ESC_PIN 5
-#define ESC_MIN 1000
-#define ESC_MAX 2000
-#define ESC_NEUTRAL 1500
-#define ESC_DEADBAND 100
-Servo motorESC;
-
-#ifdef RPM_SENSOR
-
-// Sensor for RPM
-#define SENSOR_INT_PIN 1
-#define SENSOR_DIGITAL_PIN 3
-
-#endif
-
-#define BUFF_LENGTH 5
-// variables used to read serial
-byte dummyBuff[1] = {0};
-// {start, steering, throttle, end}
-byte buffData[BUFF_LENGTH] = {0, 0, 0, 0, 0};
-bool reverseMode = false;
-
-byte expected_start = 255;
-byte expected_end = 0;
-
-// variables to detect last received serial message (safety)
-long last_received = 0;
-int maxTimout = 500;
-
-// variables to read PWM pulses
-unsigned long timer_start = 0;
-unsigned long last_interrupt_time = 0;
-long motor_speed = 0;
-long prev_motor_speed = 0;
-
-void setup()
-{
-  pinMode(LED_BUILTIN, OUTPUT);
-  digitalWrite(LED_BUILTIN, LOW);
-
-  Serial.begin(115200);
-  Serial.setTimeout(200);
-
-  // servo init
-  servoSteering.attach(SERVO_PIN, SERVO_MIN, SERVO_MAX);
-  servoSteering.writeMicroseconds(SERVO_NEUTRAL);
-
-  // motor init
-  motorESC.attach(ESC_PIN, ESC_MIN, ESC_MAX);
-  motorESC.writeMicroseconds(ESC_NEUTRAL);
-
-  #ifdef RPM_SENSOR  
-  // pinMode(SENSOR_INT_PIN, INPUT);
-  attachInterrupt(SENSOR_INT_PIN, signalChange, CHANGE);
-  #endif
-
-  #ifdef CALIBRATION_BOARD
-  // debugging board init
-  #define BUTTON_PIN 16
-  #define LED_PIN 15
-  
-  pinMode(BUTTON_PIN, INPUT);
-  pinMode(LED_PIN, OUTPUT);
-
-  // if the button is pressed within a second, enter calibration process
-  for (int i = 0; i < 20; i++)
-  {
-    delay(50);
-    if (digitalRead(BUTTON_PIN))
-    {
-      blinkLED(1);
-      waitButtonReleased(); // calibrate ESC
-      calibrationSteps();
-      break;
-    }
-  }
-  blinkLED(1);
-  #endif
-}
-
-void loop()
-{
-
-  // write rpm sensor data to the serial
-  if (Serial && motor_speed != prev_motor_speed)
-  {
-    prev_motor_speed = motor_speed;
-    Serial.println(prev_motor_speed);
-  }
-
-  if (Serial.available())
-  {
-    // read the data from the serial
-    Serial.readBytes(buffData, BUFF_LENGTH);
-
-    if (buffData[0] == expected_start && buffData[BUFF_LENGTH - 1] == expected_end) // check wether we are reading the right data buffer
-    {
-      last_received = millis();
-      changeSteering();
-      changeThrottle();
-    }
-    else
-    {
-      Serial.readBytes(dummyBuff, 1); // needed to adjust start and end of the message
-    }
-  }
-
-  // if the arduino isn't receiving anything for a given amount of time, stop the motor and servo
-  else if (millis() - last_received > maxTimout)
-  {
-    servoSteering.writeMicroseconds(SERVO_NEUTRAL);
-    motorESC.writeMicroseconds(ESC_NEUTRAL);
-  }
-}
-
-void changeSteering()
-{
-  float decoded_steering = buffData[1];
-
-  int steering = SERVO_MAX - decoded_steering / 255 * (SERVO_MAX - SERVO_MIN);
-  servoSteering.writeMicroseconds(steering);
-}
-
-void changeThrottle()
-{
-  float decoded_trottle = buffData[2];
-  float decoded_brake = buffData[3];
-
-  int throttle = ESC_MIN + decoded_trottle / 255 * (ESC_MAX - ESC_MIN);
-
-  #ifdef ENABLE_BRAKES
-  int brake = ESC_NEUTRAL - decoded_brake / 255 * (ESC_NEUTRAL - ESC_MIN);
-
-  if (brake != ESC_NEUTRAL)
-  {
-    // go back to brake mode
-    if (reverseMode)
-    {
-      motorESC.writeMicroseconds(ESC_NEUTRAL);
-    }
-    else
-    {
-      motorESC.writeMicroseconds(brake); 
-    }
-  }
-
-
-  else
-  {
-    // positive throttle
-    if (throttle >= ESC_NEUTRAL)
-    {
-      motorESC.writeMicroseconds(throttle);
-
-      if (reverseMode)
-      {
-        reverseMode = false;
-        digitalWrite(LED_BUILTIN, LOW);
-      }
-    }
-    // reverse
-    else if (reverseMode)
-    {
-      motorESC.writeMicroseconds(throttle);
-    }
-    // go into reverse mode
-    else
-    {
-      motorESC.writeMicroseconds((ESC_NEUTRAL + ESC_MIN) / 2);
-      delay(100);
-      motorESC.writeMicroseconds(ESC_NEUTRAL);
-      delay(100);
-      motorESC.writeMicroseconds(throttle); 
-      reverseMode = true;
-      digitalWrite(LED_BUILTIN, HIGH);
-    }
-  }
-  #endif
-
-  #ifndef ENABLE_BRAKES
-  motorESC.writeMicroseconds(throttle);
-  #endif
-
-}
-
-#ifdef RPM_SENSOR
-void signalChange() // this function will be called on state change of SENSOR_PIN
-{
-  last_interrupt_time = micros();
-  if (digitalRead(SENSOR_DIGITAL_PIN) == HIGH)
-  {
-    timer_start = last_interrupt_time;
-  }
-  else
-  {
-    if (timer_start != 0)
-    {
-      motor_speed = last_interrupt_time - timer_start;
-      timer_start = 0;
-    }
-  }
-}
-#endif
-
-#ifdef CALIBRATION_BOARD
-void calibrationSteps()
-{
-  // neutral pwm
-  waitButtonClicked();
-  blinkLED(1);
-  motorESC.writeMicroseconds(ESC_NEUTRAL);
-
-  // max pwm
-  waitButtonClicked();
-  blinkLED(2);
-  motorESC.writeMicroseconds(ESC_MAX);
-
-  // min pwm
-  waitButtonClicked();
-  blinkLED(3);
-  motorESC.writeMicroseconds(ESC_MIN);
-
-  delay(2900); // wait less than 3 seconds and set to neutral point to avoid motor to start at full power
-  motorESC.writeMicroseconds(ESC_NEUTRAL);
-}
-
-void waitButtonClicked()
-{
-  waitButtonPressed();
-  waitButtonReleased();
-}
-
-void waitButtonPressed()
-{
-  while (!digitalRead(BUTTON_PIN))
-  { // wait for button to be pressed
-    delay(50);
-  }
-}
-
-void waitButtonReleased()
-{
-  while (digitalRead(BUTTON_PIN))
-  { // wait for button to be released
-    delay(50);
-  }
-}
-
-void blinkLED(int rep)
-{
-  for (int i = 0; i < rep; i++)
-  {
-    digitalWrite(LED_PIN, HIGH);
-    delay(250);
-    digitalWrite(LED_PIN, LOW);
-    delay(250);
-  }
-}
-#endif
+#include <Servo.h>
+
+// #define CALIBRATION_BOARD // uncomment this line if you are using the calibration pins
+#define RPM_SENSOR // uncomment if you want to compile with the RPM_SENSOR
+// #define ENABLE_BRAKES // uncomment if you want to use the brakes
+
+// Servo
+#define SERVO_PIN 6
+#define SERVO_MIN 900
+#define SERVO_MAX 2100
+#define SERVO_NEUTRAL 1500
+#define SERVO_DEADBAND 2
+Servo servoSteering;
+
+// Motor
+#define ESC_PIN 5
+#define ESC_MIN 1000
+#define ESC_MAX 2000
+#define ESC_NEUTRAL 1500
+#define ESC_DEADBAND 100
+Servo motorESC;
+
+#ifdef RPM_SENSOR
+
+// Sensor for RPM
+#define SENSOR_INT_PIN 1
+#define SENSOR_DIGITAL_PIN 3
+
+#endif
+
+#define BUFF_LENGTH 5
+// variables used to read serial
+byte dummyBuff[1] = {0};
+// {start, steering, throttle, end}
+byte buffData[BUFF_LENGTH] = {0, 0, 0, 0, 0};
+bool reverseMode = false;
+
+byte expected_start = 255;
+byte expected_end = 0;
+
+// variables to detect last received serial message (safety)
+long last_received = 0;
+int maxTimout = 500;
+
+// variables to read PWM pulses
+unsigned long timer_start = 0;
+unsigned long last_interrupt_time = 0;
+long motor_speed = 0;
+long prev_motor_speed = 0;
+
+void setup()
+{
+  pinMode(LED_BUILTIN, OUTPUT);
+  digitalWrite(LED_BUILTIN, LOW);
+
+  Serial.begin(115200);
+  Serial.setTimeout(200);
+
+  // servo init
+  servoSteering.attach(SERVO_PIN, SERVO_MIN, SERVO_MAX);
+  servoSteering.writeMicroseconds(SERVO_NEUTRAL);
+
+  // motor init
+  motorESC.attach(ESC_PIN, ESC_MIN, ESC_MAX);
+  motorESC.writeMicroseconds(ESC_NEUTRAL);
+
+#ifdef RPM_SENSOR
+  // pinMode(SENSOR_INT_PIN, INPUT);
+  attachInterrupt(SENSOR_INT_PIN, signalChange, CHANGE);
+#endif
+
+#ifdef CALIBRATION_BOARD
+  // debugging board init
+#define BUTTON_PIN 16
+#define LED_PIN 15
+
+  pinMode(BUTTON_PIN, INPUT);
+  pinMode(LED_PIN, OUTPUT);
+
+  // if the button is pressed within a second, enter calibration process
+  for (int i = 0; i < 20; i++)
+  {
+    delay(50);
+    if (digitalRead(BUTTON_PIN))
+    {
+      blinkLED(1);
+      waitButtonReleased(); // calibrate ESC
+      calibrationSteps();
+      break;
+    }
+  }
+  blinkLED(1);
+#endif
+}
+
+void loop()
+{
+
+  // write rpm sensor data to the serial
+  if (Serial && motor_speed != prev_motor_speed)
+  {
+    prev_motor_speed = motor_speed;
+    Serial.println(prev_motor_speed);
+  }
+
+  if (Serial.available())
+  {
+    // read the data from the serial
+    Serial.readBytes(buffData, BUFF_LENGTH);
+
+    if (buffData[0] == expected_start && buffData[BUFF_LENGTH - 1] == expected_end) // check wether we are reading the right data buffer
+    {
+      last_received = millis();
+      changeSteering();
+      changeThrottle();
+    }
+    else
+    {
+      Serial.readBytes(dummyBuff, 1); // needed to adjust start and end of the message
+    }
+  }
+
+  // if the arduino isn't receiving anything for a given amount of time, stop the motor and servo
+  else if (millis() - last_received > maxTimout)
+  {
+    servoSteering.writeMicroseconds(SERVO_NEUTRAL);
+    motorESC.writeMicroseconds(ESC_NEUTRAL);
+  }
+}
+
+void changeSteering()
+{
+  float decoded_steering = buffData[1];
+
+  int steering = SERVO_MAX - decoded_steering / 255 * (SERVO_MAX - SERVO_MIN);
+  servoSteering.writeMicroseconds(steering);
+}
+
+void changeThrottle()
+{
+  float decoded_trottle = buffData[2];
+  float decoded_brake = buffData[3];
+
+  int throttle = ESC_MIN + decoded_trottle / 255 * (ESC_MAX - ESC_MIN);
+
+#ifdef ENABLE_BRAKES
+  int brake = ESC_NEUTRAL - decoded_brake / 255 * (ESC_NEUTRAL - ESC_MIN);
+
+  if (brake != ESC_NEUTRAL)
+  {
+    // go back to brake mode
+    if (reverseMode)
+    {
+      motorESC.writeMicroseconds(ESC_NEUTRAL);
+    }
+    else
+    {
+      motorESC.writeMicroseconds(brake);
+    }
+  }
+
+
+  else
+  {
+    // positive throttle
+    if (throttle >= ESC_NEUTRAL)
+    {
+      motorESC.writeMicroseconds(throttle);
+
+      if (reverseMode)
+      {
+        reverseMode = false;
+        digitalWrite(LED_BUILTIN, LOW);
+      }
+    }
+    // reverse
+    else if (reverseMode)
+    {
+      motorESC.writeMicroseconds(throttle);
+    }
+    // go into reverse mode
+    else
+    {
+      motorESC.writeMicroseconds((ESC_NEUTRAL + ESC_MIN) / 2);
+      delay(100);
+      motorESC.writeMicroseconds(ESC_NEUTRAL);
+      delay(100);
+      motorESC.writeMicroseconds(throttle);
+      reverseMode = true;
+      digitalWrite(LED_BUILTIN, HIGH);
+    }
+  }
+#endif
+
+#ifndef ENABLE_BRAKES
+  motorESC.writeMicroseconds(throttle);
+#endif
+
+}
+
+#ifdef RPM_SENSOR
+void signalChange() // this function will be called on state change of SENSOR_PIN
+{
+  last_interrupt_time = micros();
+  if (digitalRead(SENSOR_DIGITAL_PIN) == HIGH)
+  {
+    timer_start = last_interrupt_time;
+  }
+  else
+  {
+    if (timer_start != 0)
+    {
+      motor_speed = last_interrupt_time - timer_start;
+      timer_start = 0;
+    }
+  }
+}
+#endif
+
+#ifdef CALIBRATION_BOARD
+void calibrationSteps()
+{
+  // neutral pwm
+  waitButtonClicked();
+  blinkLED(1);
+  motorESC.writeMicroseconds(ESC_NEUTRAL);
+
+  // max pwm
+  waitButtonClicked();
+  blinkLED(2);
+  motorESC.writeMicroseconds(ESC_MAX);
+
+  // min pwm
+  waitButtonClicked();
+  blinkLED(3);
+  motorESC.writeMicroseconds(ESC_MIN);
+
+  delay(2900); // wait less than 3 seconds and set to neutral point to avoid motor to start at full power
+  motorESC.writeMicroseconds(ESC_NEUTRAL);
+}
+
+void waitButtonClicked()
+{
+  waitButtonPressed();
+  waitButtonReleased();
+}
+
+void waitButtonPressed()
+{
+  while (!digitalRead(BUT
+  TON_PIN))
+  { // wait for button to be pressed
+    delay(50);
+  }
+}
+
+void waitButtonReleased()
+{
+  while (digitalRead(BUTTON_PIN))
+  { // wait for button to be released
+    delay(50);
+  }
+}
+
+void blinkLED(int rep)
+{
+  for (int i = 0; i < rep; i++)
+  {
+    digitalWrite(LED_PIN, HIGH);
+    delay(250);
+    digitalWrite(LED_PIN, LOW);
+    delay(250);
+  }
+}
+#endif
diff --git a/autopylot/datasets/datagenerator.py b/autopylot/datasets/datagenerator.py
index 7cafbe5..d39154b 100644
--- a/autopylot/datasets/datagenerator.py
+++ b/autopylot/datasets/datagenerator.py
@@ -112,7 +112,10 @@ def __data_generation(self):
         for k in X.keys():
             X[k] = np.array(X[k])
         for k in Y.keys():
-            Y[k] = np.array(Y[k])
+            if "image" in k:
+                Y[k] = np.array(Y[k]) / 255.0
+            else:
+                Y[k] = np.array(Y[k])
 
         return X, Y
 
diff --git a/autopylot/datasets/transform.py b/autopylot/datasets/transform.py
index 6152cde..4efc5b1 100644
--- a/autopylot/datasets/transform.py
+++ b/autopylot/datasets/transform.py
@@ -173,44 +173,46 @@ def resize(image_data):
         )
 
     def obstacles(image_data):
-        if "image" in image_data.keys():
-            img = image_data["image"]
-
-            obstacle_path = np.random.choice(obstacles_path)
-            obstacle_img = cv2.imread(obstacle_path, cv2.IMREAD_UNCHANGED)
-
-            # need to work on the size policy
-            max_size = 32
-            upper_0 = img.shape[0] - max_size
-            upper_1 = img.shape[1] - max_size
-
-            # define random placement
-            cty = np.random.randint(min(max_size, upper_0), max(max_size, upper_0))
-            ctx = np.random.randint(min(max_size, upper_1), max(max_size, upper_1))
-            size_mult = cty / img.shape[0]
-
-            sizey = int(max_size * size_mult)
-            sizex = int(max_size * size_mult)
-            topy = cty - sizey
-            topx = ctx - sizex
-            boty = cty + sizey
-            botx = ctx + sizex
-
-            resized = cv2.resize(obstacle_img, (sizex * 2, sizey * 2))
-            color, alpha = resized[:, :, :3], resized[:, :, -1:] / 255
-
-            # apply obstacle on the image
-            img[topy:boty, topx:botx, :] = (
-                img[topy:boty, topx:botx, :] * (1 - alpha) + color * alpha
-            )
+        if "image" not in image_data.keys():
+            return 
+        
+        img = image_data["image"]
+
+        obstacle_path = np.random.choice(obstacles_path)
+        obstacle_img = cv2.imread(obstacle_path, cv2.IMREAD_UNCHANGED)
+
+        # need to work on the size policy
+        max_size = 64
+        upper_0 = img.shape[0] - max_size
+        upper_1 = img.shape[1] - max_size
+
+        # define random placement
+        cty = np.random.randint(min(max_size, upper_0), max(max_size, upper_0))
+        ctx = np.random.randint(min(max_size, upper_1), max(max_size, upper_1))
+        size_mult = cty / img.shape[0]
+
+        sizey = int(max_size * size_mult)
+        sizex = int(max_size * size_mult)
+        topy = cty - sizey
+        topx = ctx - sizex
+        boty = cty + sizey
+        botx = ctx + sizex
+
+        resized = cv2.resize(obstacle_img, (sizex * 2, sizey * 2))
+        color, alpha = resized[:, :, :3], resized[:, :, -1:] / 255
+
+        # apply obstacle on the image
+        img[topy:boty, topx:botx, :] = (
+            img[topy:boty, topx:botx, :] * (1 - alpha) + color * alpha
+        )
 
-            # car detection data
-            image_data["obstacles"] = 1  # defaults to 0
-            image_data["obstacles-size"] = size_mult  # defaults to 0
-            image_data["obstacles-coord"] = [
-                ((cty / img.shape[0]) - 0.5) * 2,
-                ((ctx / img.shape[1]) - 0.5) * 2,
-            ]  # defaults to [0, 0]
+        # car detection data
+        image_data["obstacles"] = 1  # defaults to 0
+        image_data["obstacles-size"] = size_mult  # defaults to 0
+        image_data["obstacles-coord"] = [
+            ((cty / img.shape[0]) - 0.5) * 2,
+            ((ctx / img.shape[1]) - 0.5) * 2,
+        ]  # defaults to [0, 0]
 
 
 # default values for each functions
diff --git a/autopylot/models/architectures.py b/autopylot/models/architectures.py
index 5e46084..60e7ff2 100644
--- a/autopylot/models/architectures.py
+++ b/autopylot/models/architectures.py
@@ -18,6 +18,8 @@
     Lambda,
     SeparableConv2D,
     GlobalMaxPooling2D,
+    UpSampling2D,
+    ZeroPadding2D,
 )
 from tensorflow.keras.optimizers import Adam
 
@@ -272,9 +274,9 @@ def separable_model():
 
         x = SeparableConv2D(192, 3, strides=2, use_bias=False)(x)
         x = Activation("relu")(x)
-        x = BatchNormalization()(x)
+        w = BatchNormalization()(x)
 
-        x = SeparableConv2D(256, 3, strides=1, use_bias=False)(x)
+        x = SeparableConv2D(256, 3, strides=1, use_bias=False)(w)
         x = Activation("relu")(x)
         x = BatchNormalization()(x)
 
@@ -291,12 +293,12 @@ def separable_model():
         x = Activation("relu")(x)
         x = BatchNormalization()(x)
 
-        # c1 = Dense(1, use_bias=False, activation="sigmoid", name="obstacles")(x)
-        # c2 = Dense(1, use_bias=False, activation="sigmoid", name="obstacles-size")(x)
-        # c3 = Dense(2, use_bias=False, activation="tanh", name="obstacles-coord")(x)
-        # outputs.append(c1)
-        # outputs.append(c2)
-        # outputs.append(c3)
+        c1 = Dense(1, use_bias=False, activation="sigmoid", name="obstacles")(x)
+        c2 = Dense(1, use_bias=False, activation="sigmoid", name="obstacles-size")(x)
+        c3 = Dense(2, use_bias=False, activation="tanh", name="obstacles-coord")(x)
+        outputs.append(c1)
+        outputs.append(c2)
+        outputs.append(c3)
 
         y1 = Dense(1, use_bias=False, activation="tanh", name="steering.0")(x)
         y2 = Dense(1, use_bias=False, activation="tanh", name="steering.5")(x)
@@ -307,6 +309,40 @@ def separable_model():
 
         z = Dense(3, use_bias=False, activation="softmax", name="zone")(x)
         outputs.append(z)
+        
+        # decoder
+        # (3, 8, 256)
+
+        w = SeparableConv2D(256, 3, strides=1, use_bias=False, padding="same")(w)
+        w = Activation("relu")(w)
+        w = ZeroPadding2D(padding=((1, 0), (1, 0)))(w)
+        w = UpSampling2D(size=(2, 2))(w)
+        # (8, 18, 256)
+
+        w = SeparableConv2D(192, 3, strides=1, use_bias=False, padding="same")(w)
+        w = Activation("relu")(w)
+        w = ZeroPadding2D(padding=((1, 1), (1, 1)))(w)
+        w = UpSampling2D(size=(2, 2))(w)
+        # (20, 40, 192)
+
+        w = SeparableConv2D(96, 5, strides=1, use_bias=False, padding="same")(w)
+        w = Activation("relu")(w)
+        w = UpSampling2D(size=(2, 2))(w)
+        # (40, 80, 96)
+
+        w = SeparableConv2D(48, 5, strides=1, use_bias=False, padding="same")(w)
+        w = Activation("relu")(w)
+        w = UpSampling2D(size=(2, 2))(w)
+        # (80, 160, 48)
+
+        w = SeparableConv2D(24, 5, strides=1, use_bias=False, padding="same")(w)
+        w = Activation("relu")(w)
+
+        w = Conv2D(3, 3, strides=1, use_bias=False, padding="same")(w)
+        w = Activation("tanh")(w)
+        w = ZeroPadding2D(padding=((20, 20), (0, 0)))(w)
+        w = Lambda(lambda x: (x + 1) * 0.5, name="image.0")(w)
+        outputs.append(w)
 
         # Create the model
         model = Model(inputs=inputs, outputs=outputs)
@@ -314,8 +350,8 @@ def separable_model():
         # Compile it
         model.compile(
             optimizer=Adam(),
-            loss=["mse", "mse", "mse", "mse", "mse", "mse", "categorical_crossentropy"],
-            loss_weights=[1, 1, 1, 1, 1, 1, 0.75],
+            loss=["mse", "mse", "mse", "mse", "mse", "mse", "categorical_crossentropy", "mse"],
+            loss_weights=[1, 1, 1, 1, 1, 1, 0.75, 1],
         )
 
         logging.info(f"created separable model with {get_flops(model)} FLOPS")
diff --git a/autopylot/models/train.py b/autopylot/models/train.py
index 346457c..e460f7e 100644
--- a/autopylot/models/train.py
+++ b/autopylot/models/train.py
@@ -62,7 +62,6 @@ def train(
         train_split=settings.TRAIN_SPLITS,
         shuffle=settings.TRAIN_SHUFFLE,
         verbose=settings.TRAIN_VERBOSE,
-        do_save=True,
         additionnal_funcs=[],
     ):
         """Trains the model on the given dataset.
@@ -135,7 +134,7 @@ def train(
             validation_data=test_generator,
             validation_steps=1,
             epochs=settings.TRAIN_EPOCHS,
-            workers=4,
+            workers=8,
         )
 
         if settings.MODEL_SAVE_SETTINGS:
@@ -147,4 +146,10 @@ def train(
                 "shuffle": shuffle,
                 "verbose": verbose,
             }
+        
+        # save alternative version without decoder
+        encoder = utils.remove_decoder(self.model)
+        if encoder is not None:
+            utils.save_model(encoder, f"{self.name}_enc", model_info=self.model_info)
+
         utils.save_model(self.model, self.name, model_info=self.model_info)
diff --git a/autopylot/models/utils.py b/autopylot/models/utils.py
index ac37cd3..ffd5515 100644
--- a/autopylot/models/utils.py
+++ b/autopylot/models/utils.py
@@ -144,11 +144,38 @@ def get_clean_layer_name(name):
     Returns:
         string: the filtered layer name.
     """
+    
+    if name.endswith("_output"):
+        name = name.split("_output")[0]
+
     delimit_chars = ":_/"
     for char in delimit_chars:
         name = name.split(char)[0]
+
     return name
 
+def remove_decoder(model):
+    """Remove the decoder from the model.
+    removes the "image" output layer and rebuild the model.
+
+    Args:
+        model (Model): the model.
+
+    Returns:
+        Model: the model without the decoder. OR None if no changes
+    """
+
+    new_outputs = []
+    outputs = model.outputs
+    for output in outputs:
+        if "image" not in output.name:
+            new_outputs.append(output)
+
+    if len(new_outputs) == 0:
+        raise ValueError("No output left after removing decoder")
+    
+    if len(new_outputs) != len(outputs):
+        return Model(model.inputs, new_outputs)
 
 def predict_decorator(func, inputs, outputs):
     """Decorate the model.predict function.
diff --git a/autopylot/tests/test_io.py b/autopylot/tests/test_io.py
index d6751d2..a761d1c 100644
--- a/autopylot/tests/test_io.py
+++ b/autopylot/tests/test_io.py
@@ -34,45 +34,52 @@ def test_create_directory():
 
 def test_load_image_none():
     """Testing if the image was loaded."""
-    image = io.load_image(os.getcwd() + "\\testing_io\\test.png")
-    assert image is None, "should not be None, the image doens't exist."
+    path = os.path.join(os.getcwd(), "testing_io", "test.png")
+    image = io.load_image(path)
+    assert image is None, "should be None, the image doens't exist."
 
 
 def test_save_image():
     """Testing the saving of an image."""
     image = np.zeros((2, 2, 3), dtype=np.uint8)
-    save = io.save_image(os.getcwd() + "\\testing_io\\test.png", image)
+    path = os.path.join(os.getcwd(), "testing_io", "test.png")
+    save = io.save_image(path, image)
     assert save is True, "Image not saved."
 
 
 def test_load_image():
     """Testing the loading of the image."""
-    image = io.load_image(os.getcwd() + "\\testing_io\\test.png")
+    path = os.path.join(os.getcwd(), "testing_io", "test.png")
+    image = io.load_image(path)
     assert image.shape == (2, 2, 3)
 
 
 def test_load_json_none():
     """Testing the loading of a non existing json (should raise an error)."""
     with pytest.raises(Exception):
-        io.load_json(os.getcwd() + "\\testing_io\\test.json")
+        path = os.path.join(os.getcwd(), "testing_io", "test.json")
+        io.load_json(path)
 
 
 def test_save_json():
     """Testing the saving of a dictionnary into a .json file."""
     data = {"test": "this is a test"}
-    save = io.save_json(os.getcwd() + "\\testing_io\\test.json", data)
+    path = os.path.join(os.getcwd(), "testing_io", "test.json")
+    save = io.save_json(path, data)
     assert save is True, "json not saved."
 
 
 def test_load_json():
     """Testing the loading of the .json file."""
-    data = io.load_json(os.getcwd() + "\\testing_io\\test.json")
+    path = os.path.join(os.getcwd(), "testing_io", "test.json")
+    data = io.load_json(path)
     assert data == {"test": "this is a test"}
 
 
 def test_load_image_data():
     """Testing the loading of both image and .json file."""
-    image_data = io.load_image_data(os.getcwd() + "\\testing_io\\test.json")
+    path = os.path.join(os.getcwd(), "testing_io", "test.json")
+    image_data = io.load_image_data(path)
     image = image_data["image"]
     del image_data["image"]
     assert image.shape == (2, 2, 3) and image_data == {"test": "this is a test"}
@@ -85,8 +92,9 @@ def test_save_image_data():
         "image": np.zeros((2, 2, 3), dtype=np.uint8),
     }
 
-    io.save_image_data(image_data, os.getcwd() + "\\testing_io\\test2.json")
-    image_data_copy = io.load_image_data(os.getcwd() + "\\testing_io\\test2.json")
+    path = os.path.join(os.getcwd(), "testing_io", "test2.json")
+    io.save_image_data(image_data, path)
+    image_data_copy = io.load_image_data(path)
 
     assert (
         image_data["image"].shape == image_data_copy["image"].shape
diff --git a/autopylot/utils/vis.py b/autopylot/utils/vis.py
index 02856b1..b4e6d71 100644
--- a/autopylot/utils/vis.py
+++ b/autopylot/utils/vis.py
@@ -41,7 +41,7 @@ def vis_line_scalar(
     return vis_image
 
 
-def vis_steering(image_data, image_key="image"):
+def vis_steering(image_data, image_key="image", color=(0, 0, 255)):
     """Visualize the steering scalar.
 
     Args:
@@ -57,12 +57,12 @@ def vis_steering(image_data, image_key="image"):
         pos=(0.5, 0),
         length=(0, 0.25),
         fact=(0.2, 0),
-        color=(0, 0, 255),
+        color=color,
         thickness=2,
     )
 
 
-def vis_throttle(image_data, image_key="image"):
+def vis_throttle(image_data, image_key="image", color=(0, 0, 255)):
     """Visualize the throttle scalar.
 
     Args:
@@ -73,7 +73,6 @@ def vis_throttle(image_data, image_key="image"):
         np.array: modified image with the drawn visualization.
     """
     throttle = image_data["throttle"]
-    color = (0, 0, 255)
     if throttle < 0:
         color = (255, 0, 0)
         throttle *= -1
@@ -215,7 +214,7 @@ def vis_obstacles(image_data, image_key="image", x_minmax=(-1, 1), y_minmax=(-1,
     return image_data[image_key]
 
 
-def vis_all(image_data, image_key="image"):
+def vis_all(image_data, image_key="image", color=(0, 0, 255)):
     """Visualize every data present in the image_data dictionary.
 
     Args:
diff --git a/main_programs/examples/load_and_vis_data.py b/main_programs/examples/load_and_vis_data.py
index 1142a73..bf9ff8c 100644
--- a/main_programs/examples/load_and_vis_data.py
+++ b/main_programs/examples/load_and_vis_data.py
@@ -24,7 +24,10 @@
 
 
 def main(path):
-    for path in dataset.sort_paths(dataset.get_every_json_paths(path)):
+    paths = dataset.sort_paths(dataset.get_every_json_paths(path))
+    i = 0 
+    while i < len(paths):
+        path = paths[i]
         image_data = io.load_image_data(path)
         transformer(image_data)
         vis_image = vis.vis_all(image_data)
@@ -39,7 +42,20 @@ def main(path):
 
         vis.cv2.imshow("vis_image", vis_image)
         vis.cv2.imshow("vis_pred", vis_pred)
-        vis.cv2.waitKey(0)
+
+        if "image.0" in image_data:
+            vis.cv2.imshow("image.0", image_data["image.0"])
+
+        key = vis.cv2.waitKey(0)
+        if key == ord("q"):
+            vis.cv2.destroyAllWindows()
+            break
+        elif key == ord("n"):
+            i += 300
+        elif key == ord("p"):
+            i -= 300
+        else:
+            i += 1
 
 
 if __name__ == "__main__":
diff --git a/main_programs/setup.py b/main_programs/setup.py
deleted file mode 100644
index 4875c1e..0000000
--- a/main_programs/setup.py
+++ /dev/null
@@ -1,13 +0,0 @@
-"""
-Before working with autopylot, which include training and predicting,
-you need to setup the environment first. To do so, just run the following
-command:
-
-    python setup.py
-
-!!! Make sure autopylot is installed !!!
-
-it will generate a settings.json file in autopylot's root directory and
-it will create $HOME/dataset and $HOME/collect directories.
-"""
-from autopylot.utils import settings
diff --git a/main_programs/train.py b/main_programs/train.py
index fb04286..eaa530f 100644
--- a/main_programs/train.py
+++ b/main_programs/train.py
@@ -24,6 +24,5 @@
         batch_size=settings.TRAIN_BATCH_SIZE,
         train_split=settings.TRAIN_SPLITS,
         verbose=settings.TRAIN_VERBOSE,
-        do_save=True,
         additionnal_funcs=[],
     )
diff --git a/models/auto_encoder_test/auto_encoder_test.h5 b/models/auto_encoder_test/auto_encoder_test.h5
new file mode 100644
index 0000000..c034922
Binary files /dev/null and b/models/auto_encoder_test/auto_encoder_test.h5 differ
diff --git a/models/auto_encoder_test/auto_encoder_test.info b/models/auto_encoder_test/auto_encoder_test.info
new file mode 100644
index 0000000..e56202b
--- /dev/null
+++ b/models/auto_encoder_test/auto_encoder_test.info
@@ -0,0 +1 @@
+{"inputs": [["image", [120, 160, 3]]], "outputs": [["steering.0", [1]], ["steering.5", [1]], ["steering.10", [1]], ["zone", [3]], ["image.0", [120, 160, 3]]], "settings": {"dataset_path": "/home/maxime/datasets", "epochs": 1, "batch_size": 64, "train_split": 0.9, "shuffle": true, "verbose": 1}}
\ No newline at end of file
diff --git a/models/auto_encoder_test/auto_encoder_test.tflite b/models/auto_encoder_test/auto_encoder_test.tflite
new file mode 100644
index 0000000..d7afa3e
Binary files /dev/null and b/models/auto_encoder_test/auto_encoder_test.tflite differ
diff --git a/models/auto_encoder_test_enc/auto_encoder_test_enc.h5 b/models/auto_encoder_test_enc/auto_encoder_test_enc.h5
new file mode 100644
index 0000000..0719979
Binary files /dev/null and b/models/auto_encoder_test_enc/auto_encoder_test_enc.h5 differ
diff --git a/models/auto_encoder_test_enc/auto_encoder_test_enc.info b/models/auto_encoder_test_enc/auto_encoder_test_enc.info
new file mode 100644
index 0000000..e56202b
--- /dev/null
+++ b/models/auto_encoder_test_enc/auto_encoder_test_enc.info
@@ -0,0 +1 @@
+{"inputs": [["image", [120, 160, 3]]], "outputs": [["steering.0", [1]], ["steering.5", [1]], ["steering.10", [1]], ["zone", [3]], ["image.0", [120, 160, 3]]], "settings": {"dataset_path": "/home/maxime/datasets", "epochs": 1, "batch_size": 64, "train_split": 0.9, "shuffle": true, "verbose": 1}}
\ No newline at end of file
diff --git a/models/auto_encoder_test_enc/auto_encoder_test_enc.tflite b/models/auto_encoder_test_enc/auto_encoder_test_enc.tflite
new file mode 100644
index 0000000..633cfd8
Binary files /dev/null and b/models/auto_encoder_test_enc/auto_encoder_test_enc.tflite differ
diff --git a/models/auto_encoder_zones/auto_encoder_zones.h5 b/models/auto_encoder_zones/auto_encoder_zones.h5
new file mode 100644
index 0000000..284ebdd
Binary files /dev/null and b/models/auto_encoder_zones/auto_encoder_zones.h5 differ
diff --git a/models/auto_encoder_zones/auto_encoder_zones.info b/models/auto_encoder_zones/auto_encoder_zones.info
new file mode 100644
index 0000000..4664744
--- /dev/null
+++ b/models/auto_encoder_zones/auto_encoder_zones.info
@@ -0,0 +1 @@
+{"inputs": [["image", [120, 160, 3]]], "outputs": [["obstacles", [1]], ["obstacles-size", [1]], ["obstacles-coord", [2]], ["steering.0", [1]], ["steering.5", [1]], ["steering.10", [1]], ["zone", [3]], ["image.0", [120, 160, 3]]], "settings": {"dataset_path": "/home/maxime/datasets", "epochs": 1, "batch_size": 64, "train_split": 0.9, "shuffle": true, "verbose": 1}}
\ No newline at end of file
diff --git a/models/auto_encoder_zones/auto_encoder_zones.tflite b/models/auto_encoder_zones/auto_encoder_zones.tflite
new file mode 100644
index 0000000..d484bf5
Binary files /dev/null and b/models/auto_encoder_zones/auto_encoder_zones.tflite differ
diff --git a/models/auto_encoder_zones_enc/auto_encoder_zones_enc.h5 b/models/auto_encoder_zones_enc/auto_encoder_zones_enc.h5
new file mode 100644
index 0000000..8f4d9d9
Binary files /dev/null and b/models/auto_encoder_zones_enc/auto_encoder_zones_enc.h5 differ
diff --git a/models/auto_encoder_zones_enc/auto_encoder_zones_enc.info b/models/auto_encoder_zones_enc/auto_encoder_zones_enc.info
new file mode 100644
index 0000000..4664744
--- /dev/null
+++ b/models/auto_encoder_zones_enc/auto_encoder_zones_enc.info
@@ -0,0 +1 @@
+{"inputs": [["image", [120, 160, 3]]], "outputs": [["obstacles", [1]], ["obstacles-size", [1]], ["obstacles-coord", [2]], ["steering.0", [1]], ["steering.5", [1]], ["steering.10", [1]], ["zone", [3]], ["image.0", [120, 160, 3]]], "settings": {"dataset_path": "/home/maxime/datasets", "epochs": 1, "batch_size": 64, "train_split": 0.9, "shuffle": true, "verbose": 1}}
\ No newline at end of file
diff --git a/models/auto_encoder_zones_enc/auto_encoder_zones_enc.tflite b/models/auto_encoder_zones_enc/auto_encoder_zones_enc.tflite
new file mode 100644
index 0000000..09c660e
Binary files /dev/null and b/models/auto_encoder_zones_enc/auto_encoder_zones_enc.tflite differ
diff --git a/setup.py b/setup.py
index 565f0a7..c47eb0c 100644
--- a/setup.py
+++ b/setup.py
@@ -1,3 +1,16 @@
+"""
+Before working with autopylot, which include training and predicting,
+you need to setup the environment first. To do so, just run the following
+command:
+
+    python setup.py
+
+!!! Make sure autopylot is installed !!!
+
+it will generate a settings.json file in autopylot's root directory and
+it will create $HOME/dataset and $HOME/collect directories.
+"""
+
 from setuptools import setup
 
 setup(
@@ -17,7 +30,7 @@
         "glob2",
         "pyserial",
         "python-socketio[client]",
-        "gym-donkeycar @ git+https://github.com/tawnkramer/gym-donkeycar",
+        "gym-donkeycar @ git+https://github.com/Autonomobile/gym-donkeycar",
         "keyboard",
         "keras-flops",
         "numpy-quaternion",