Repository: bark-simulator/bark-ml (language: python)

Issue:
Implementation and Usage of the DeterministicScenarioGeneration
This enables the tests to be more meaningful and easier debugging capabilities for reinforcement learning algorithms.


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index dafd3ed9..097455d2 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -13,9 +13,15 @@ jobs:
       - name: Setting up virtual environment
         run: virtualenv -p python3 ./bark_ml/python_wrapper/venv  --system-site-packages
       - name: Entering the virtual python environment..
-        run: . ./bark_ml/python_wrapper/venv/bin/activate
-      - name: Installing additional Python packages..
+        run: . ./bark_ml/python_wrapper/venv/bin/activate 
+      - name: Installing Tensorflow packages..
         run: pip install tensorflow-addons==0.11.2 graph-nets==1.1.0 ray==1.0.0
+      - name: Installing Pytorch venv..
+        run: pip install --target=./bark_ml/python_wrapper/venv/lib/python3.7/site-packages torch==1.9.0+cpu torchvision -f https://download.pytorch.org/whl/torch_stable.html
+      - name: Installing Pytorch global..
+        run: pip install torch==1.9.0+cpu torchvision -f https://download.pytorch.org/whl/torch_stable.html
+      - name: Add torch environment variable
+        run: export LD_LIBRARY_PATH=./bark_ml/python_wrapper/venv/lib/python3.7/site-packages/torch/lib/
       - name: Running tests
         run: bazel test //bark_ml:unit_tests
 
diff --git a/.gitignore b/.gitignore
index e96676ca..d186fe70 100644
--- a/.gitignore
+++ b/.gitignore
@@ -166,8 +166,13 @@ tests/data/database_maps_DR_DEU_Merging_MT_shifted.xodr
 *.png
 *.p
 
+/checkpoints**
+/summaries**
+
+*events.out*
+*.pt*
 /experiments/configs/**
 /experiments/configs/**/**
 /experiments/runs/**/**
 /experiments/runs/**/**
-/_build/**
\ No newline at end of file
+/_build/**
diff --git a/.vscode/launch.json b/.vscode/launch.json
index 1eb991f5..f5ee1973 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -11,8 +11,9 @@
 					"program": "${file}",
 					"console": "integratedTerminal",
 					"env": {
-						"PYTHONPATH": "${workspaceFolder}/bazel-bark-ml/examples/tfa_gnn_simple/"
-			 }
+						"PYTHONPATH": "${workspaceFolder}/bazel-bark-ml/bark_ml/commons/",
+						"LD_LIBRARY_PATH" : "$LD_LIBRARY_PATH:${workspaceFolder}/bark_ml/python_wrapper/venv/lib/python3.7/site-packages/torch/lib/"
+			 		}
 			},
 			{
 					"name": "(gdb) Launch",
diff --git a/WORKSPACE b/WORKSPACE
index 139c7cc3..9d3c8503 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -1,6 +1,8 @@
 workspace(name = "bark_ml")
 load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
 
+load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
+
 load("//utils:dependencies.bzl", "bark_ml_dependencies")
 bark_ml_dependencies()
 
@@ -16,7 +18,7 @@ python_configure(name = "local_config_python")
 # -------- Benchmark Database -----------------------
 git_repository(
   name = "benchmark_database",
-  commit="ff6e433ecb7878ebe59996f3994ff67483a7c297",
+  commit="119864d1940110074237386861ab6e4ce8d04477",
   remote = "https://github.com/bark-simulator/benchmark-database"
 )
 
@@ -24,4 +26,3 @@ load("@benchmark_database//util:deps.bzl", "benchmark_database_dependencies")
 load("@benchmark_database//load:load.bzl", "benchmark_database_release")
 benchmark_database_dependencies()
 benchmark_database_release()
-# --------------------------------------------------
diff --git a/bark_ml/BUILD b/bark_ml/BUILD
index f3c60b81..1d5bd776 100644
--- a/bark_ml/BUILD
+++ b/bark_ml/BUILD
@@ -28,7 +28,7 @@ test_suite(
     "//bark_ml/tests:py_library_tfa_tests",
     "//bark_ml/tests:py_bark_behavior_model_tests",
     "//bark_ml/experiment:py_experiment_tests",
-    "//bark_ml/tests:py_apollo_tests"
+    "//bark_ml/tests:py_apollo_tests",
   ]
 )
 
diff --git a/bark_ml/behaviors/discrete_behavior.py b/bark_ml/behaviors/discrete_behavior.py
index f5d6a493..836785c5 100644
--- a/bark_ml/behaviors/discrete_behavior.py
+++ b/bark_ml/behaviors/discrete_behavior.py
@@ -50,6 +50,13 @@ def __init__(self,
     default_model = BehaviorMacroActionsFromParamServer(self._macro_action_params)
     super().__init__(self._macro_action_params, default_model.GetMotionPrimitives())
 
+  def __setstate__(self, macro_params):
+    default_model = BehaviorMacroActionsFromParamServer(macro_params)
+    super().__init__(macro_params, default_model.GetMotionPrimitives())
+  
+  def __getstate__(self):
+    return self._macro_action_params
+
   @property
   def action_space(self):
     return Discrete(len(self.GetMotionPrimitives()))
\ No newline at end of file
diff --git a/bark_ml/commons/debug_settings.py b/bark_ml/commons/debug_settings.py
index 1e7cda83..cf8192b1 100755
--- a/bark_ml/commons/debug_settings.py
+++ b/bark_ml/commons/debug_settings.py
@@ -3,11 +3,11 @@
 
 cwd = os.getcwd()
 workspace_folder = cwd
-repo_paths = ["bark_project", "bark_project/python", "benchmark_database", "com_github_interaction_dataset_interaction_dataset", \
+repo_paths = ["bark_project", "bark_project/python_wrapper", "benchmark_database", "com_github_interaction_dataset_interaction_dataset", \
         "com_github_interaction_dataset_interaction_dataset/python", "diadem_project", "bark_ml"]
 
 executed_file = sys.argv[0]
-tmp = executed_file.replace("examples", "bazel-bin/examples")
+tmp = executed_file.replace("bark_ml", "bazel-bin/bark_ml")
 runfiles_dir = tmp.replace(".py", ".runfiles")
 
 sys.path.append(runfiles_dir)
diff --git a/bark_ml/evaluators/base_evaluator.hpp b/bark_ml/evaluators/base_evaluator.hpp
index a7a3df15..f2f422aa 100644
--- a/bark_ml/evaluators/base_evaluator.hpp
+++ b/bark_ml/evaluators/base_evaluator.hpp
@@ -48,7 +48,7 @@ using bark::models::dynamic::StateDefinition::X_POSITION;
 using bark::models::dynamic::StateDefinition::Y_POSITION;
 using bark::models::dynamic::StateDefinition::THETA_POSITION;
 using bark::models::dynamic::StateDefinition::VEL_POSITION;
-using ObservedState = Eigen::Matrix<float, 1, Eigen::Dynamic>;
+using ObservedState = Eigen::Matrix<double, 1, Eigen::Dynamic>;
 using bark::commons::transformation::FrenetPosition;
 using State = Eigen::Matrix<double, Eigen::Dynamic, 1>;
 
diff --git a/bark_ml/evaluators/goal_reached.hpp b/bark_ml/evaluators/goal_reached.hpp
index 2bbeabf9..7912f361 100644
--- a/bark_ml/evaluators/goal_reached.hpp
+++ b/bark_ml/evaluators/goal_reached.hpp
@@ -52,7 +52,7 @@ using bark::models::dynamic::StateDefinition::X_POSITION;
 using bark::models::dynamic::StateDefinition::Y_POSITION;
 using bark::models::dynamic::StateDefinition::THETA_POSITION;
 using bark::models::dynamic::StateDefinition::VEL_POSITION;
-using ObservedState = Eigen::Matrix<float, 1, Eigen::Dynamic>;
+using ObservedState = Eigen::Matrix<double, 1, Eigen::Dynamic>;
 using bark::commons::transformation::FrenetPosition;
 using State = Eigen::Matrix<double, Eigen::Dynamic, 1>;
 using bark_ml::evaluators::BaseEvaluator;
diff --git a/bark_ml/library_wrappers/lib_diadem/BUILD b/bark_ml/library_wrappers/lib_diadem/BUILD
deleted file mode 100644
index 164a753f..00000000
--- a/bark_ml/library_wrappers/lib_diadem/BUILD
+++ /dev/null
@@ -1,11 +0,0 @@
-py_library(
-    name = "diadem",
-    srcs = ["diadem_bark_environment.py"],
-    data = ["@bark_project//bark/python_wrapper:core.so"],
-    imports=["diadem_project"],
-    deps = [
-        "@bark_project//bark/runtime:runtime",
-        "@diadem_project//diadem:__init__"
-        ],
-    visibility = ["//visibility:public"],
-)
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_diadem/diadem_bark_environment.py b/bark_ml/library_wrappers/lib_diadem/diadem_bark_environment.py
deleted file mode 100644
index 5206c0ff..00000000
--- a/bark_ml/library_wrappers/lib_diadem/diadem_bark_environment.py
+++ /dev/null
@@ -1,53 +0,0 @@
-# Copyright (c) 2020 The diadem authors
-#
-# This work is licensed under the terms of the MIT license.
-# For a copy, see <https://opensource.org/licenses/MIT>.
-# ========================================================
-#
-
-from diadem.environments import Environment
-
-class DiademBarkEnvironment(Environment):
-    """BARK environment for the diadem library."""
-
-    def __init__(self, runtime,  params=None):
-        super().__init__()
-        self.runtime = runtime
-
-    def _initialize_params(self, params=None):
-        super()._initialize_params(params=params)
-
-    def step(self, actions):
-        next_observation, reward, done, info = self.runtime.step(actions)
-        return next_observation, reward, done, info
-
-    def reset(self, state=None):
-        current_observation = self.runtime.reset()
-        return current_observation, False
-
-    def create_init_states(self, size=None, idx=None):
-        pass
-
-    def render(self, filename=None, show=True):
-        self.runtime.render()
-
-    @property
-    def actionspace(self):
-        return self.runtime.action_space
-
-    @property
-    def observationspace(self):
-        return self.runtime.observation_space
-
-    @staticmethod
-    def contains_training_data():
-        return True
-
-    @staticmethod
-    def state_as_string():
-        return "BarkEnvironment" #todo: maybe print scenario idx
-
-    @property
-    def action_timedelta(self):
-        # only for visualization purposes actual step time unknown
-        return self.runtime.step_time
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/BUILD b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/BUILD
index d717516f..c09ab40d 100644
--- a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/BUILD
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/BUILD
@@ -6,6 +6,19 @@ py_library(
         ":fqf_agent",
         ":iqn_agent",
         ":qrdqn_agent",
+        ":imitation_agent",
+        ":carin_agent",
+        "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent:util",
+        "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations:demonstrations",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+
+py_library(
+    name = "util",
+    srcs = [
+        "util.py"
     ],
     visibility = ["//visibility:public"],
 )
@@ -13,11 +26,26 @@ py_library(
 
 py_library(
     name = "base_agent",
-    srcs = ["base_agent.py"],
-    data = ["@bark_project//bark/python_wrapper:core.so"],
+    srcs = ["base_agent.py"
+    ],
+    data = [
+        "@bark_project//bark:generate_core"
+        ],
     deps = [
         "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/memory:init",
+        "//bark_ml/behaviors:behaviors",
+    ],
+)
+
+py_library(
+    name = "imitation_agent",
+    srcs = ["imitation_agent.py"],
+    deps = [":base_agent",
+            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/loss:loss_function",
+            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model:model",
+            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn:utils",
     ],
+    visibility = ["//visibility:public"],
 )
 
 
@@ -51,3 +79,9 @@ py_library(
     visibility = ["//visibility:public"],
 )
 
+py_library(
+    name = "carin_agent",
+    srcs = ["carin_agent.py"],
+    deps = [":imitation_agent"],
+    visibility = ["//visibility:public"],
+)
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/__init__.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/__init__.py
index 6543b935..12a85d2a 100644
--- a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/__init__.py
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/__init__.py
@@ -1,3 +1,8 @@
-from .fqf_agent import FQFAgent  # pylint: disable=unused-import
-from .iqn_agent import IQNAgent  # pylint: disable=unused-import
-from .qrdqn_agent import QRDQNAgent  # pylint: disable=unused-import
+from .fqf_agent import FQFAgent
+from .iqn_agent import IQNAgent
+from .qrdqn_agent import QRDQNAgent
+from .imitation_agent import ImitationAgent, PolicyImitationAgent
+from .carin_agent import CarinAgent
+from .base_agent import BaseAgent, TrainingBenchmark
+from .demonstrations import *
+from .util import *
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/base_agent.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/base_agent.py
index 2bd16c47..363a04b0 100644
--- a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/base_agent.py
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/base_agent.py
@@ -12,6 +12,7 @@
 from torch.utils.tensorboard import SummaryWriter
 import torch
 import numpy as np
+import pickle
 import os
 from abc import abstractmethod
 
@@ -24,95 +25,220 @@
 from bark.core.models.behavior import BehaviorModel
 import logging
 
+def to_pickle(obj, dir, file):
+  path = os.path.join(dir, file)
+  with open(path, 'wb') as handle:
+    pickle.dump(obj, handle, protocol=pickle.HIGHEST_PROTOCOL)
+
+def from_pickle(dir, file):
+  path = os.path.join(dir, file)
+  with open(path, 'rb') as handle:
+    obj = pickle.load(handle)
+  return obj
+
+class TrainingBenchmark:
+  def __init__(self):
+    self.training_env = None
+    self.num_episodes = None
+    self.max_episode_steps = None
+    self.agent = None
+
+  def reset(self, training_env, num_episodes, max_episode_steps, agent):
+    self.training_env = training_env
+    self.num_episodes = num_episodes
+    self.max_episode_steps = max_episode_steps
+    self.agent = agent
 
-class BaseAgent(BehaviorModel):
-  """BaseAgent that can be used in BARK and BARK-ML."""
+  def run(self):
+    # returns dict with evaluated metrics
+    num_episodes = 0
+    total_return = 0.0
+
+    while True:
+      state = self.training_env.reset()
+      episode_steps = 0
+      episode_return = 0.0
+      done = False
+      while (not done) and episode_steps <= self.max_episode_steps:
+        if self.agent.is_random(eval=True):
+          action = self.agent.explore()
+        else:
+          action = self.agent.Act(state)
+
+        next_state, reward, done, _ = self.training_env.step(action)
+        episode_steps += 1
+        episode_return += reward
+        state = next_state
+
+      num_episodes += 1
+      total_return += episode_return
+
+      if num_episodes > self.num_episodes:
+        break
+
+    mean_return = total_return / num_episodes
+    return {"mean_return" : mean_return}, f"Mean return: {mean_return}"
 
-  def __init__(self, env, test_env, params, bark_behavior=None):
+  def is_better(self, eval_result1, than_eval_result2):
+    return eval_result1["mean_return"] > than_eval_result2["mean_return"]
+
+
+
+
+class BaseAgent(BehaviorModel):
+  def __init__(self, agent_save_dir=None, env=None, params=None, training_benchmark=None, checkpoint_load=None, 
+              demonstrations = None):
     BehaviorModel.__init__(self, params)
     self._params = params
-    self.env = env
-    self.test_env = test_env
-    self._bark_behavior_model = bark_behavior or BehaviorDiscreteMacroActionsML(params)
-
-    self.device = torch.device("cuda" if self._params["ML"]["BaseAgent"][
-        "Cuda", "", True] and torch.cuda.is_available() else "cpu")
+    self._env = env
+    self._training_benchmark = training_benchmark or TrainingBenchmark()
+    self._agent_save_dir = agent_save_dir
+    self._demonstrations = demonstrations
+    self._checkpoint_load = checkpoint_load
+    
+    if (not checkpoint_load) and params is not None:
+      self.reset_params(self._params["ML"]["BaseAgent"])
+      self.reset_action_observer(env)
+      self.init_always()
+      self.reset_training_variables()
+    elif checkpoint_load:
+      self.load_pickable_members(agent_save_dir)
+      self.load_other()
+      self.init_always()
+      self.load_models(BaseAgent.check_point_directory(agent_save_dir, checkpoint_load))
+    else:
+      raise ValueError("Unusual param combination for agent initialization.")
 
-    self.online_net = None
-    self.target_net = None
 
-    self.steps = 0
-    self.learning_steps = 0
-    self.episodes = 0
-    self.best_eval_score = -np.inf
-    self.num_actions = self.env.action_space.n
-    self.num_steps = self._params["ML"]["BaseAgent"]["NumSteps", "", 5000000]
-    self.batch_size = self._params["ML"]["BaseAgent"]["BatchSize", "", 32]
-
-    self.double_q_learning = self._params["ML"]["BaseAgent"][
-        "Double_q_learning", "", False]
-    self.dueling_net = self._params["ML"]["BaseAgent"]["DuelingNet", "", False]
-    self.noisy_net = self._params["ML"]["BaseAgent"]["NoisyNet", "", False]
-    self.use_per = self._params["ML"]["BaseAgent"]["Use_per", "", False]
-
-    self.reward_log_interval = self._params["ML"]["BaseAgent"]["RewardLogInterval", "", 5]
-    self.summary_log_interval = self._params["ML"]["BaseAgent"]["SummaryLogInterval", "", 100]
-    self.eval_interval = self._params["ML"]["BaseAgent"]["EvalInterval", "",
-                                                         25000]
-    self.num_eval_steps = self._params["ML"]["BaseAgent"]["NumEvalSteps", "",
-                                                          12500]
-    self.gamma_n = \
-     self._params["ML"]["BaseAgent"]["Gamma", "", 0.99] ** \
-     self._params["ML"]["BaseAgent"]["Multi_step", "", 1]
+  def init_always(self):
+    self.device = torch.device("cuda" if (self.use_cuda and torch.cuda.is_available()) else "cpu")
 
-    self.start_steps = self._params["ML"]["BaseAgent"]["StartSteps", "", 5000]
-    self.epsilon_train = LinearAnneaer(
-        1.0, self._params["ML"]["BaseAgent"]["EpsilonTrain", "", 0.01],
-        self._params["ML"]["BaseAgent"]["EpsilonDecaySteps", "", 25000])
-    self.epsilon_eval = self._params["ML"]["BaseAgent"]["EpsilonEval", "",
-                                                        0.001]
-    self.update_interval = \
-     self._params["ML"]["BaseAgent"]["Update_interval", "", 4]
-    self.target_update_interval = self._params["ML"]["BaseAgent"][
-        "TargetUpdateInterval", "", 5000]
-    self.max_episode_steps = \
-     self._params["ML"]["BaseAgent"]["MaxEpisodeSteps",  "", 10000]
-    self.grad_cliping = self._params["ML"]["BaseAgent"]["GradCliping", "", 5.0]
-
-    self.summary_dir = \
-     self._params["ML"]["BaseAgent"]["SummaryPath", "", ""]
-    self.model_dir = \
-     self._params["ML"]["BaseAgent"]["CheckpointPath", "", ""]
-
-    if not os.path.exists(self.model_dir) and self.model_dir:
-      os.makedirs(self.model_dir)
-    if not os.path.exists(self.summary_dir) and self.summary_dir:
-      os.makedirs(self.summary_dir)
-
-    self.writer = SummaryWriter(log_dir=self.summary_dir)
+    logging.info(f"Summary writer directory: {os.path.abspath(BaseAgent.summary_dir(self.agent_save_dir))}")
+    self.writer = SummaryWriter(log_dir=BaseAgent.summary_dir(self.agent_save_dir))
     self.train_return = RunningMeanStats(self.summary_log_interval)
 
+    if not os.path.exists(BaseAgent.summary_dir(self.agent_save_dir)):
+      os.makedirs(BaseAgent.summary_dir(self.agent_save_dir))
+
     # NOTE: by default we do not want the action to be set externally
     #       as this enables the agents to be plug and played in BARK.
     self._set_action_externally = False
+    self._training_benchmark.reset(self._env, \
+        self.num_eval_episodes, self.max_episode_steps, self)
+
+  def reset_action_observer(self, env):
+    self._observer = self._env._observer
+    self._ml_behavior = self._env._ml_behavior
+
+  def clean_pickables(self, pickables):
+    del pickables["online_net"]
+    del pickables["_agent_save_dir"]
+    del pickables["target_net"]
+    del pickables["_env"]
+    del pickables["_training_benchmark"]
+    del pickables["_checkpoint_load"]
+    del pickables["device"]
+    del pickables["writer"]
+
+
+  def save_pickable_members(self, pickable_dir):
+    if not os.path.exists(pickable_dir):
+      os.makedirs(pickable_dir)
+    pickables = dict(self.__dict__)
+    self.clean_pickables(pickables)
+    to_pickle(pickables, pickable_dir, "agent_pickables")
+
+  def load_pickable_members(self, agent_save_dir):
+    pickables = from_pickle(BaseAgent.pickable_directory(agent_save_dir), "agent_pickables")
+    try:
+      del pickables["_agent_save_dir"]
+      del pickables["base_demonstrations_dir"]
+    except:
+      pass
+    self.__dict__.update(pickables)
+
+  def load_other(self):
+    pass
 
+  def reset_training_variables(self):
     # Replay memory which is memory-efficient to store stacked frames.
     if self.use_per:
       beta_steps = (self.num_steps - self.start_steps) / \
              self.update_interval
       self.memory = LazyPrioritizedMultiStepMemory(
-          self._params["ML"]["BaseAgent"]["MemorySize", "", 10**6],
-          self.env.observation_space.shape,
+          self.memory_size,
+          self.observer.observation_space.shape,
           self.device,
-          self._params["ML"]["BaseAgent"]["Gamma", "", 0.99],
-          self._params["ML"]["BaseAgent"]["Multi_step", "", 1],
+          self.gamma,
+          self.multi_step,
           beta_steps=beta_steps)
     else:
       self.memory = LazyMultiStepMemory(
-          self._params["ML"]["BaseAgent"]["MemorySize", "", 10**6],
-          self.env.observation_space.shape, self.device,
-          self._params["ML"]["BaseAgent"]["Gamma", "", 0.99],
-          self._params["ML"]["BaseAgent"]["Multi_step", "", 1])
+          self.memory_size,
+          self.observer.observation_space.shape,
+          self.device,
+          self.gamma,
+          self.multi_step)
+
+    self.steps = 0
+    self.learning_steps = 0
+    self.episodes = 0
+    self.best_eval_results = None
+
+  def reset_params(self, params):
+    self.num_steps = params["NumSteps", "", 5000000]
+    self.batch_size = params["BatchSize", "", 32]
+
+    self.double_q_learning = params["Double_q_learning", "", False]
+    self.dueling_net = params["DuelingNet", "", False]
+    self.noisy_net = params["NoisyNet", "", False]
+    self.use_per = params["Use_per", "", False]
+
+    self.reward_log_interval = params["RewardLogInterval", "", 5]
+    self.summary_log_interval = params["SummaryLogInterval", "", 100]
+    self.eval_interval = params["EvalInterval", "", 25000]
+    self.num_eval_episodes = params["NumEvalEpisodes", "",
+                                                          12500]
+    self.gamma_n = params["Gamma", "", 0.99] ** \
+        params["Multi_step", "", 1]
+
+    self.start_steps = params["StartSteps", "", 5000]
+    self.epsilon_train = LinearAnneaer(
+        1.0, params["EpsilonTrain", "", 0.01],
+        params["EpsilonDecaySteps", "", 25000])
+    self.epsilon_eval = params["EpsilonEval", "",
+                                                        0.001]
+    self.update_interval = params["Update_interval", "", 4]
+    self.target_update_interval = params["TargetUpdateInterval", "", 5000]
+    self.max_episode_steps = params["MaxEpisodeSteps",  "", 10000]
+    self.grad_cliping = params["GradCliping", "", 5.0]
+
+    self.memory_size = params["MemorySize", "", 10**6]
+    self.gamma = params["Gamma", "", 0.99]
+    self.multi_step = params["Multi_step", "", 1]
+
+    self.use_cuda = params["Cuda", "", False] 
+
+  @property
+  def observer(self):
+      return self._observer
+
+  @property 
+  def env(self):
+    return self._env
+
+  @property
+  def ml_behavior(self):
+    return self._ml_behavior
+
+  @property
+  def num_actions(self):
+    return self.ml_behavior.action_space.n
+  
+  @property
+  def agent_save_dir(self):
+    return self._agent_save_dir
+
 
   def run(self):
     while True:
@@ -140,7 +266,7 @@ def update_target(self):
 
   def explore(self):
     # Act with randomness.
-    action = self.env.action_space.sample()
+    action = self.ml_behavior.action_space.sample()
     return action
 
   @property
@@ -171,16 +297,18 @@ def calculate_actions(self, state):
   def Plan(self, dt, observed_world):
     # NOTE: if training is enabled the action is set externally
     if not self._set_action_externally:
-      observed_state = self.env._observer.Observe(observed_world)
+      observed_state = self.observer.Observe(observed_world)
       action = self.Act(observed_state)
       self._action = action
 
     action = self._action
     # set action to be executed
-    self._bark_behavior_model.ActionToBehavior(action)
-    trajectory = self._bark_behavior_model.Plan(dt, observed_world)
+    self._ml_behavior.ActionToBehavior(action)
+    trajectory = self._ml_behavior.Plan(dt, observed_world)
+    dynamic_action = self._ml_behavior.GetLastAction()
     # NOTE: BARK requires models to have trajectories of the past
     BehaviorModel.SetLastTrajectory(self, trajectory)
+    BehaviorModel.SetLastAction(self, dynamic_action)
     return trajectory
 
   @abstractmethod
@@ -192,25 +320,61 @@ def Clone(self):
 
   @property
   def action_space(self):
-    return self._bark_behavior_model.action_space
+    return self._ml_behavior.action_space
+
+  @staticmethod
+  def check_point_directory(agent_save_dir, checkpoint_type):
+    return os.path.join(agent_save_dir, "checkpoints/", checkpoint_type)
+
+  @staticmethod
+  def pickable_directory(agent_save_dir):
+    return os.path.join(agent_save_dir, "pickable/")
 
-  def save_models(self, save_dir):
-    if not os.path.exists(save_dir):
-      os.makedirs(save_dir)
+  @staticmethod
+  def summary_dir(agent_save_dir):
+    return os.path.join(agent_save_dir, "summaries")
+
+  def save_models(self, checkpoint_dir):
+    if not os.path.exists(checkpoint_dir):
+      os.makedirs(checkpoint_dir)
     torch.save(self.online_net.state_dict(),
-               os.path.join(save_dir, 'online_net.pth'))
+               os.path.join(checkpoint_dir, 'online_net.pth'))
     torch.save(self.target_net.state_dict(),
-               os.path.join(save_dir, 'target_net.pth'))
+               os.path.join(checkpoint_dir, 'target_net.pth'))
     online_net_script = torch.jit.script(self.online_net)
-    online_net_script.save(os.path.join(save_dir, 'online_net_script.pt'))
-
-  def load_models(self, save_dir):
-    self.online_net.load_state_dict(
-        torch.load(os.path.join(save_dir, 'online_net.pth')))
-    self.target_net.load_state_dict(
-        torch.load(os.path.join(save_dir, 'target_net.pth')))
+    online_net_script.save(os.path.join(checkpoint_dir, 'online_net_script.pt'))
+
+  def get_script_filename(self, checkpoint_load=None):
+    if not checkpoint_load:
+      checkpoint_load = self._checkpoint_load
+    checkpoint_dir = BaseAgent.check_point_directory(self._agent_save_dir, checkpoint_load)
+    return os.path.join(checkpoint_dir, 'online_net_script.pt')
+
+  def save_in_dir(self, agent_save_dir, checkpoint_type):
+    self._agent_save_dir = agent_save_dir
+    self.save_models(BaseAgent.check_point_directory(agent_save_dir, checkpoint_type))
+    self.save_pickable_members(BaseAgent.pickable_directory(agent_save_dir))
+
+  def save(self, checkpoint_type="last"):
+    self.save_in_dir(self.agent_save_dir, checkpoint_type)
+
+  def load_models(self, checkpoint_dir):
+    try: 
+      self.online_net.load_state_dict(
+        torch.load(os.path.join(checkpoint_dir, 'online_net.pth')))
+    except RuntimeError:
+      self.online_net.load_state_dict(
+        torch.load(os.path.join(checkpoint_dir, 'online_net.pth'), map_location=torch.device('cpu')))
+    try: 
+      self.target_net.load_state_dict(
+        torch.load(os.path.join(checkpoint_dir, 'target_net.pth')))
+    except RuntimeError:
+      self.target_net.load_state_dict(
+        torch.load(os.path.join(checkpoint_dir, 'target_net.pth'), map_location=torch.device('cpu')))
 
   def visualize(self, num_episodes=5):
+    if not self.env:
+      raise ValueError("No environment available for visualization. Was agent reloaded?")
     for _ in range(0, num_episodes):
       state = self.env.reset()
       done = False
@@ -280,48 +444,33 @@ def train_step_interval(self):
 
     if self.steps % self.eval_interval == 0:
       self.evaluate()
-      self.save_models(os.path.join(self.model_dir, 'final'))
+      self.save("final")
       self.online_net.train()
 
-  def evaluate(self):
+  def raw_evaluation_score(self):
     self.online_net.eval()
-    num_episodes = 0
-    num_steps = 0
-    total_return = 0.0
+    eval_results, _ = self._training_benchmark.run()
+    return eval_results
 
-    while True:
-      state = self.test_env.reset()
-      episode_steps = 0
-      episode_return = 0.0
-      done = False
-      while (not done) and episode_steps <= self.max_episode_steps:
-        if self.is_random(eval=True):
-          action = self.explore()
-        else:
-          action = self.Act(state)
 
-        next_state, reward, done, _ = self.test_env.step(action)
-        num_steps += 1
-        episode_steps += 1
-        episode_return += reward
-        state = next_state
-
-      num_episodes += 1
-      total_return += episode_return
-
-      if num_steps > self.num_eval_steps:
-        break
-
-    mean_return = total_return / num_episodes
+  def evaluate(self):
+    if not self._training_benchmark:
+      logging.info("No evaluation performed since no training benchmark available.")
+      return
+    self.online_net.eval()
+    
+    eval_results, formatted_result = self._training_benchmark.run()
 
-    if mean_return > self.best_eval_score:
-      self.best_eval_score = mean_return
-      self.save_models(os.path.join(self.model_dir, 'best'))
+    if not self.best_eval_results or \
+        self._training_benchmark.is_better(eval_results, self.best_eval_results):
+      self.best_eval_results = eval_results
+      self.save("best")
 
     # We log evaluation results along with training frames = 4 * steps.
-    self.writer.add_scalar('return/test', mean_return, 4 * self.steps)
+    for eval_result_name, eval_result in eval_results.items():
+      self.writer.add_scalar(eval_result_name, eval_result, 4 * self.steps)
     logging.info('-' * 60)
-    logging.info(f'Num steps: {self.steps:<5}  ' f'return: {mean_return:<5.1f}')
+    logging.info('Evaluation result: {}'.format(formatted_result))
     logging.info('-' * 60)
 
   def __del__(self):
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/carin_agent.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/carin_agent.py
new file mode 100644
index 00000000..1e0030d0
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/carin_agent.py
@@ -0,0 +1,184 @@
+from collections import OrderedDict
+import torch
+from torch import nn, sigmoid, cat, unsqueeze, flatten
+from torch.optim import Adam, RMSprop
+
+from .imitation_agent import ImitationAgent
+from bark_ml.core.value_converters import *
+
+def init_weights(m):
+  if isinstance(m, nn.Linear):
+    nn.init.xavier_uniform_(m.weight)
+    m.bias.data.fill_(0.01)
+  elif isinstance(m, nn.Conv1d):
+    nn.init.xavier_uniform_(m.weight)
+    m.bias.data.fill_(0.01)
+
+
+class CarinAgent(ImitationAgent):
+  def __init__(self, *args, **kwargs):
+    super(CarinAgent, self).__init__(*args, **kwargs)
+
+  def init_network(self):
+    self.online_net = Carin(
+        num_channels=self.observer.observation_space.shape[0],
+        num_actions=self.num_actions,
+        num_value_functions=self.num_value_functions,
+        params=self._params).to(self.device)
+
+    self.optim = RMSprop(self.online_net.parameters(),
+                         lr=self.learning_rate,
+                         weight_decay=self.weight_decay,
+                         alpha=0.95,
+                         eps=0.00001)
+
+
+class Carin(nn.Module):
+  """
+  CARIN = Custom ARchitecture ImitatioN
+
+  Usage of CNN layers inspired by:
+    'Combining Planning and Deep Reinforcement Learning in Tactical
+    Decision Making for Autonomous Driving' (Hoel et al., 2019)
+  """
+  def __init__(self, num_channels, num_actions, num_value_functions, params):
+    super(Carin, self).__init__()
+    self.num_channels = num_channels
+    self.num_actions = num_actions
+    self.num_value_functions = num_value_functions
+
+    self.num_features_ego = params["ML"]["CarinModel"]["NumFeaturesEgo", "", 6]
+    self.num_features_other_agent = params["ML"]["CarinModel"][
+        "NumFeaturesOtherAgent", "", 4]
+    self.num_other_agents = \
+      (self.num_channels - self.num_features_ego) // self.num_features_other_agent
+
+    self.dropout_p = params["ML"]["ImitationModel"]["DropoutProbability", "",
+                                                    0]
+
+    self.input_conv_dims = params["ML"]["CarinModel"]["InputConvDims", "", []]
+    self.shared_layer_dims = params["ML"]["ImitationModel"]["EmbeddingDims",
+                                                            "",
+                                                            [256, 256, 256]]
+    self.head_networks_dims = params["ML"]["CarinModel"][
+        "MultitaskLearningEmbeddingDims", "", []]
+
+    # CNN layers (transform input features of each agent separately).
+    input_conv_layers, last_input_layer_dim = self.make_input_conv_layers()
+    self.input_conv_net = nn.Sequential(input_conv_layers)
+    self.input_conv_net.apply(init_weights)
+
+    # Shared FC layers
+    shared_layers, last_shared_layer_dim = self.make_shared_layers(
+        last_input_layer_dim)
+    self.shared_network = nn.Sequential(shared_layers)
+    self.shared_network.apply(init_weights)
+
+    # Multitask learning (separate FC networks for each value function)
+    head_layers = self.make_head_layers(last_shared_layer_dim)
+    self.head_networks = nn.ModuleList()
+    for hl in head_layers:
+      net = nn.Sequential(hl)
+      net.apply(init_weights)
+      self.head_networks.append(net)
+    
+    self.value_converter = NNToValueConverterSequential(self.num_actions)
+
+  @torch.jit.unused
+  @property
+  def nn_to_value_converter(self):
+    return self.value_converter
+
+  def forward(self, states):
+    # Treat ego features and features of other agents separately
+    ego_features = states[:, :self.num_features_ego]
+    other_agent_features = states[:, self.num_features_ego:]
+
+    # Convert input tensor from 2D to 3D (needed for convolution afterwards)
+    other_agent_features = unsqueeze(other_agent_features, 1)
+
+    # Use CNN to transform features of each agent separately
+    other_agent_features = self.input_conv_net(other_agent_features)
+
+    # Transform back to 2D
+    other_agent_features = flatten(other_agent_features,
+                                   start_dim=1,
+                                   end_dim=2)
+
+    # Concatenate ego features with the features of other agents
+    transformed_input = cat([ego_features, other_agent_features], 1)
+
+    # Fully connected layers
+    shared_features = self.shared_network(transformed_input)
+
+    # Compute output for each value function, then concat them
+    outputs = [head(shared_features) for head in self.head_networks]
+    action_values = cat(outputs, 1)
+
+    if not self.training:
+      # Evaluation phase, output values between 0 and 1
+      action_values = sigmoid(action_values)
+    return action_values
+
+  def make_input_conv_layers(self):
+    tuple_list = []
+    in_channels = 1
+    kernel_size = self.num_features_other_agent
+    stride = kernel_size
+    for idx, layer in enumerate(self.input_conv_dims):
+      out_channels = layer
+      tuple_list.append((f"conv_layer{idx}",
+                         nn.Conv1d(in_channels,
+                                   out_channels,
+                                   kernel_size=kernel_size,
+                                   stride=stride)))
+      tuple_list.append((f"relu{idx}", nn.ReLU()))
+      in_channels = out_channels
+      kernel_size = 1
+      stride = 1
+
+    if len(self.input_conv_dims) > 0:
+      tuple_list.append(
+          ("maxpool", nn.MaxPool2d(kernel_size=(1, self.num_other_agents))))
+      # Ego features are appended to the end of the output of CNN
+      last_dimension = in_channels + self.num_features_ego
+    else:
+      last_dimension = self.num_channels
+
+    return OrderedDict(tuple_list), last_dimension
+
+  def make_shared_layers(self, start_layer_dim):
+    """
+    Creates layers that are shared among all predictions
+    """
+    tuple_list = []
+    last_dim = start_layer_dim
+    for idx, layer in enumerate(self.shared_layer_dims):
+      current_dim = layer
+      tuple_list.append((f"layer{idx}", nn.Linear(last_dim, current_dim)))
+      tuple_list.append((f"relu{idx}", nn.ReLU()))
+      if self.dropout_p != 0:
+        tuple_list.append((f"dropout{idx}", nn.Dropout(p=self.dropout_p)))
+      last_dim = current_dim
+    return OrderedDict(tuple_list), last_dim
+
+  def make_head_layers(self, start_layer_dim):
+    """
+    Creates separate layers for each value function output (multitask learning)
+    """
+    head_layers = []
+    for i in range(1, self.num_value_functions + 1):
+      tuple_list = []
+      last_dim = start_layer_dim
+      for idx, layer in enumerate(self.head_networks_dims):
+        current_dim = layer
+        tuple_list.append(
+            (f"head{i}_layer{idx}", nn.Linear(last_dim, current_dim)))
+        tuple_list.append((f"head{i}_relu{idx}", nn.ReLU()))
+        if self.dropout_p != 0:
+          tuple_list.append(
+              (f"head{i}_dropout{idx}", nn.Dropout(p=self.dropout_p)))
+        last_dim = current_dim
+      tuple_list.append((f"output{i}", nn.Linear(last_dim, self.num_actions)))
+      head_layers.append(OrderedDict(tuple_list))
+    return head_layers
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations/BUILD b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations/BUILD
new file mode 100644
index 00000000..d3d656a7
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations/BUILD
@@ -0,0 +1,13 @@
+py_library(
+    name = "demonstrations",
+    srcs = ["demonstration_collector.py",
+            "__init__.py"
+            ],
+    deps = [
+            "@bark_project//bark/benchmark:benchmark_runner",
+            "@bark_project//bark/benchmark:benchmark_runner_mp",
+            "@bark_project//bark/benchmark:benchmark_result",
+    ],
+    visibility = ["//visibility:public"],
+)
+
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations/__init__.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations/__init__.py
new file mode 100644
index 00000000..0c4a2542
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations/__init__.py
@@ -0,0 +1,3 @@
+
+from .demonstration_collector import DemonstrationCollector
+from .demonstration_collector import ActionValuesCollector
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations/demonstration_collector.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations/demonstration_collector.py
new file mode 100644
index 00000000..242b40fb
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations/demonstration_collector.py
@@ -0,0 +1,345 @@
+# Copyright (c) 2020 fortiss GmbH
+#
+# Authors: Julian Bernhard, Klemens Esterle, Patrick Hart and
+# Tobias Kessler
+#
+# This software is released under the MIT License.
+# https://opensource.org/licenses/MIT
+
+import logging
+import pickle
+import os
+import numpy as np
+
+from bark.runtime.commons.parameters import ParameterServer
+
+from bark.core.world.evaluation import *
+from bark.core.world import *
+
+from bark.benchmark.benchmark_result import BenchmarkResult
+from bark.benchmark.benchmark_runner import BenchmarkRunner, BehaviorConfig
+from bark.benchmark.benchmark_runner_mp import BenchmarkRunnerMP
+
+from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.agent.util import *
+
+
+def to_pickle(obj, dir, file):
+  path = os.path.join(dir, file)
+  with open(path, 'wb') as handle:
+    pickle.dump(obj, handle, protocol=pickle.HIGHEST_PROTOCOL)
+
+def from_pickle(dir, file):
+  path = os.path.join(dir, file)
+  with open(path, 'rb') as handle:
+    obj = pickle.load(handle)
+  return obj
+
+
+class DemonstrationEvaluator(BaseEvaluator):
+  def __init__(self, observer, reward_evaluator):
+    super(DemonstrationEvaluator, self).__init__()
+    self._observer = observer
+    self._reward_evaluator = reward_evaluator
+    self._agent_id = None
+    self._last_nn_state = None
+    self._episode_experiences = []
+
+  def SetAgentId(self, agent_id):
+    self._agent_id = agent_id
+
+  def GetNNInputState(self, observed_world):
+    return self._observer.Observe(observed_world)
+
+  def GetAction(self, observed_world):
+    return observed_world.agents[self._agent_id].behavior_model.GetLastMacroAction()
+
+  def GetStepEvaluation(self, observed_world, action):
+    return self._reward_evaluator.Evaluate(observed_world, action)
+
+  def Evaluate(self, world):
+    if isinstance(world, World):
+      experience, done, info = self.GetExperience(world)
+      self._episode_experiences.append(experience)
+      return (self._episode_experiences, done, info)
+    else:
+      raise NotImplementedError()
+
+  def MakeExperienceTuple(self, nn_state, action, reward, next_nn_state, done):
+    demo = True
+    return (nn_state, action, reward, next_nn_state, done, demo)
+
+  def GetExperience(self, world):
+    if len(self._episode_experiences) == 0:
+      self._reward_evaluator.Reset(world)
+    observed_world = world.Observe([self._agent_id])[0]
+    current_nn_state = self.GetNNInputState(observed_world)
+    action = self.GetAction(observed_world)
+    reward, done, info = self.GetStepEvaluation(observed_world, action)
+    experience = self.MakeExperienceTuple(self._last_nn_state, action, reward, current_nn_state, done)
+    self._last_nn_state = current_nn_state
+    return experience, done, info
+
+  def __setstate__(self, d):
+    self._observer = d["observer"]
+    self._reward_evaluator = d["reward_evaluator"]
+    self._agent_id = None
+    self._last_nn_state = None
+    self._episode_experiences = []
+
+  def __getstate__(self):
+    return {"observer" : self._observer, \
+          "reward_evaluator" : self._reward_evaluator}
+
+class ActionValueEvaluator(BaseEvaluator):
+  def __init__(self, observer):
+    super(ActionValueEvaluator, self).__init__()
+    self._agent_id = None
+    self._observer = observer
+    self._episode_experiences = []
+
+  def SetAgentId(self, agent_id):
+    self._agent_id = agent_id
+
+  def GetExperience(self, world):
+    observed_world = world.Observe([self._agent_id])[0]
+    current_nn_state = self._observer.Observe(observed_world)
+    action_values = self.GetActionValues(observed_world)
+    policy = self.GetPolicy(observed_world)
+    return current_nn_state, action_values, policy
+
+  def AddMissingActionsValues(self, value_dict, num_actions):
+    values = []
+    assert(num_actions >= len(value_dict))
+    for idx in  range(0, num_actions):
+      if idx in value_dict:
+        values.append(value_dict[idx])
+      else:
+        values.append(np.nan)
+    return values
+
+  def GetPolicy(self, observed_world):
+    behavior = observed_world.agents[self._agent_id].behavior_model
+    num_actions = len(behavior.ego_behavior.GetMotionPrimitives())
+    policy = behavior.last_policy_sampled[1]
+    return self.AddMissingActionsValues(policy, num_actions)
+
+  def GetActionValues(self, observed_world):
+    behavior = observed_world.agents[self._agent_id].behavior_model
+    num_actions = len(behavior.ego_behavior.GetMotionPrimitives())
+    cost_envelope_values = []
+    cost_collision_values = []
+    if "envelope" in behavior.last_cost_values:
+      cost_envelope_values = \
+           self.AddMissingActionsValues(behavior.last_cost_values["envelope"], num_actions)
+    cost_collision_values = []
+    if "collision" in behavior.last_cost_values:
+      cost_collision_values = \
+        self.AddMissingActionsValues(behavior.last_cost_values["collision"], num_actions)
+    return_values = self.AddMissingActionsValues(behavior.last_return_values, num_actions)
+    action_values = []
+    action_values.extend(cost_envelope_values)
+    action_values.extend(cost_collision_values)
+    action_values.extend(return_values)
+    return action_values
+
+  def Evaluate(self, world):
+    if isinstance(world, World):
+      experience = self.GetExperience(world)
+      self._episode_experiences.append(experience)
+      return self._episode_experiences
+    else:
+      raise NotImplementedError()
+
+  def __setstate__(self, d):
+    self._observer = d["observer"]
+    self._agent_id = None
+    self._episode_experiences = []
+
+  def __getstate__(self):
+    return {"observer" : self._observer}
+
+
+class DemonstrationCollector:
+  def __init__(self):
+    self._collection_result = None
+    self._demonstrations = None
+    self._directory = None
+    self._observer = None
+    self._motion_primitive_behavior = None
+
+  def _GetDefaultRunnerInitParams(self):
+    return {"log_eval_avg_every" : 5}
+
+  def _GetDefaultRunnerRunParams(self):
+    return {"maintain_history" : False, "checkpoint_every" : 10, "viewer" : None}
+
+  def CollectDemonstrations(self, num_episodes, directory, motion_primitive_behavior, 
+       env=None, observer=None, reward_evaluator=None, benchmark_configs=None,
+       use_mp_runner=True, runner_init_params = None,
+      runner_run_params=None):
+
+    if env:
+      observer = env._observer
+      reward_evaluator = env._evaluator
+      scenario_generator = env._scenario_generator
+    else:
+      scenario_generator = None
+    if env:
+      behaviors = {"demo_behavior" : motion_primitive_behavior}
+    else:
+      behaviors = None
+    self._observer = observer
+    self._motion_primitive_behavior = motion_primitive_behavior
+    demo_evaluator = self.GetEvaluators(observer, reward_evaluator)
+    evaluators = {**default_training_evaluators(), "demo_evaluator" : demo_evaluator}
+    terminal_when = self.GetTerminalCriteria()
+
+    runner_init_params_def = self._GetDefaultRunnerInitParams()
+    runner_init_params_def.update(runner_init_params or {})
+    runner_type = BenchmarkRunnerMP if use_mp_runner else BenchmarkRunner
+    runner = runner_type(evaluators=evaluators,
+                                  scenario_generation=scenario_generator,
+                                  benchmark_configs = benchmark_configs,
+                                  terminal_when=terminal_when,
+                                  behaviors=behaviors,
+                                  num_scenarios = num_episodes,
+                                  **runner_init_params_def)
+    runner.clear_checkpoint_dir()
+    runner_run_params_def = self._GetDefaultRunnerRunParams()
+    runner_run_params_def.update(runner_run_params or {})
+    self._collection_result = runner.run(**runner_run_params_def)
+    self.dump(directory)
+    return self._collection_result
+
+  def GetEvaluators(self, observer, reward_evaluator):
+    return DemonstrationEvaluator(observer, reward_evaluator)
+
+  def GetTerminalCriteria(self, *args):
+    return {"demo_evaluator" : lambda x : x[1] == True} # second index in evaluation result is done
+
+  def dump(self, directory):
+    self._directory = directory
+    if not os.path.exists(directory):
+      os.makedirs(directory)
+    to_pickle(self._observer, directory, DemonstrationCollector.observer_filename())
+    to_pickle(self._motion_primitive_behavior, directory, DemonstrationCollector.motion_primitive_behavior_filename())
+    if self._collection_result:
+      self._collection_result.dump(os.path.join(directory, DemonstrationCollector.collection_result_filename()), dump_histories=False, dump_configs=False)
+    if self._demonstrations:
+      to_pickle(self._demonstrations, directory, DemonstrationCollector.demonstrations_filename())
+
+  def GetDirectory(self):
+    return self._directory
+
+  @staticmethod
+  def _load(collector, directory):
+    collection_result_fullname = os.path.join(directory, DemonstrationCollector.collection_result_filename())
+    if os.path.exists(collection_result_fullname):
+      collector._collection_result = BenchmarkResult.load(collection_result_fullname)
+    else:
+      logging.warning("Collection result not existing.")
+    demonstration_fullname = os.path.join(directory, DemonstrationCollector.demonstrations_filename())
+    if os.path.exists(demonstration_fullname):
+      collector._demonstrations = from_pickle(directory, DemonstrationCollector.demonstrations_filename())
+    else:
+      logging.warning("Demonstrations not existing.")
+    collector._directory = directory
+    collector._observer = from_pickle(directory, DemonstrationCollector.observer_filename())
+    collector._motion_primitive_behavior = from_pickle(directory, DemonstrationCollector.motion_primitive_behavior_filename())
+    return collector
+
+  @classmethod
+  def load(cls, directory):
+    collector = DemonstrationCollector()
+    return DemonstrationCollector._load(collector, directory)
+
+  @property
+  def observer(self):
+    return self._observer
+
+  @property
+  def motion_primitive_behavior(self):
+    return self._motion_primitive_behavior
+
+  @staticmethod
+  def collection_result_filename():
+    return "collection_result"
+
+  @staticmethod
+  def demonstrations_filename():
+    return "demonstrations"
+
+  @staticmethod
+  def observer_filename():
+    return "observer"
+
+  @staticmethod
+  def motion_primitive_behavior_filename():
+    return "motion_primitive_behavior"
+
+  def UseCollectedRow(self, row, eval_criteria):
+    demo_eval_result, done, info = row["demo_evaluator"]
+    use_scenario = True
+    for crit, func in eval_criteria.items():
+      use_scenario = use_scenario and func(info[crit])
+    return use_scenario
+
+  def GetDemonstrations(self, row):
+    demo_eval_result, done, info = row["demo_evaluator"]
+    return demo_eval_result[1:]
+
+  def ProcessCollectionResult(self, eval_criteria = None):
+    if not self._collection_result:
+      logging.error("Collection results not created yet. Call CollectDemonstrations first.")
+      return
+
+    data_frame = self._collection_result.get_data_frame()
+    self._demonstrations = []
+    exceptions = data_frame[data_frame.Terminal == "exception_raised"]
+    if (len(exceptions.index) > 0):
+      logging.warning(f"Removing {len(exceptions.index)} with raised exceptions")
+      data_frame = data_frame[~(data_frame.Terminal == "exception_raised")]
+    for index, row in data_frame.iterrows():
+      if not eval_criteria or self.UseCollectedRow(row, eval_criteria):
+        self._demonstrations.extend(self.GetDemonstrations(row))
+    self.dump(self._directory)
+    return self._demonstrations
+
+  def GetDemonstrationExperiences(self):
+    if not self._demonstrations:
+      self.ProcessCollectionResult()
+    return self._demonstrations
+
+  def GetCollectionResult(self):
+    return self._collection_result
+  
+
+class ActionValuesCollector(DemonstrationCollector):
+  def __init__(self, terminal_criteria):
+    super(ActionValuesCollector, self).__init__()
+    self.terminal_criteria = terminal_criteria
+
+  def UseCollectedRow(self, row, eval_criteria):
+    use_scenario = True
+    for crit, func in eval_criteria.items():
+      use_scenario = use_scenario and func(row[crit])
+    return use_scenario
+
+  def GetDemonstrations(self, row):
+    demos = None
+    try:
+      demos = [[list(tp[0]), list(tp[1]), list(tp[2])] for tp in list(row["demo_evaluator"][1:])]
+    except:
+      demos = [[list(tp[0]), list(tp[1])] for tp in list(row["demo_evaluator"][1:])]
+    return demos
+
+  def GetEvaluators(self, observer, reward_evaluator):
+    return ActionValueEvaluator(observer)
+
+  def GetTerminalCriteria(self, *args):
+    return self.terminal_criteria
+
+  @classmethod
+  def load(cls, directory):
+    collector = ActionValuesCollector(terminal_criteria=None)
+    return DemonstrationCollector._load(collector, directory)
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/fqf_agent.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/fqf_agent.py
index ddf0a2bf..7dfbfa5a 100644
--- a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/fqf_agent.py
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/fqf_agent.py
@@ -20,11 +20,11 @@
 
 
 class FQFAgent(BaseAgent):
-  """FQFAgent that can be used in BARK and BARK-ML."""
-
-  def __init__(self, env, test_env, params):
-    super(FQFAgent, self).__init__(env, test_env, params)
+  def __init__(self, *args, **kwargs):
+    super(FQFAgent, self).__init__(*args, **kwargs)
 
+  def reset_params(self, params):
+    super(FQFAgent, self).reset_params(params)
     # NOTE: The author said the training of Fraction Proposal Net is
     # unstable and value distribution degenerates into a deterministic
     # one rarely (e.g. 1 out of 20 seeds). So you can use entropy of value
@@ -34,9 +34,13 @@ def __init__(self, env, test_env, params):
     self.N = self._params["ML"]["FQFAgent"]["N", "", 32]
     self.num_cosines = self._params["ML"]["FQFAgent"]["NumCosines", "", 64]
     self.kappa = self._params["ML"]["FQFAgent"]["Kappa", "", 1.0]
-
-    # Online network.
-    self.online_net = FQF(num_channels=env.observation_space.shape[0],
+    self.fractional_learning_rate = self._params["ML"]["FQFAgent"]["FractionalLearningRate", "",
+                                          2.5e-9]
+  
+  def init_always(self):
+    super(FQFAgent, self).init_always()
+      # Online network.
+    self.online_net = FQF(num_channels=self.observer.observation_space.shape[0],
                           num_actions=self.num_actions,
                           N=self.N,
                           num_cosines=self.num_cosines,
@@ -44,7 +48,7 @@ def __init__(self, env, test_env, params):
                           noisy_net=self.noisy_net,
                           params=self._params).to(self.device)
     # Target network.
-    self.target_net = FQF(num_channels=env.observation_space.shape[0],
+    self.target_net = FQF(num_channels=self.observer.observation_space.shape[0],
                           num_actions=self.num_actions,
                           N=self.N,
                           num_cosines=self.num_cosines,
@@ -60,8 +64,7 @@ def __init__(self, env, test_env, params):
 
     self.fraction_optim = RMSprop(
         self.online_net.fraction_net.parameters(),
-        lr=self._params["ML"]["FQFAgent"]["FractionalLearningRate", "",
-                                          2.5e-9],
+        lr=self.fractional_learning_rate,
         alpha=0.95,
         eps=0.00001)
 
@@ -72,6 +75,11 @@ def __init__(self, env, test_env, params):
         lr=self._params["ML"]["FQFAgent"]["QuantileLearningRate", "", 5e-5],
         eps=1e-2 / self.batch_size)
 
+  def clean_pickables(self, pickables):
+    super(FQFAgent, self).clean_pickables(pickables)
+    del pickables["fraction_optim"]
+    del pickables["quantile_optim"]
+
   def update_target(self):
     self.target_net.dqn_net.load_state_dict(
         self.online_net.dqn_net.state_dict())
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/imitation_agent.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/imitation_agent.py
new file mode 100644
index 00000000..f81f3c7c
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/imitation_agent.py
@@ -0,0 +1,439 @@
+# Copyright (c) 2020 fortiss GmbH
+#
+# Authors: Patrick Hart, Julian Bernhard, Klemens Esterle, and
+# Tobias Kessler, Mansoor Nasir
+#
+# This software is released under the MIT License.
+# https://opensource.org/licenses/MIT
+
+from collections import deque
+import numpy as np
+import logging
+import os
+import math
+import random
+
+import torch
+from torch import nn
+from torch.optim import Adam, RMSprop, AdamW
+
+from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.model import Imitation, PolicyImitation
+from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.agent.demonstrations import ActionValuesCollector
+from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.utils \
+ import disable_gradients, update_params, \
+ calculate_quantile_huber_loss, evaluate_quantile_at_action
+from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.agent.loss.loss_function \
+  import apply_sigmoid_to_dict, LossMSE, LossBCE, LossPolicyCrossEntropy, \
+  LossHuber, LossTukey, LossEpsInsensitiveHuber, LossRelative
+from .base_agent import BaseAgent, TrainingBenchmark
+
+class BenchmarkSupervisedLoss(TrainingBenchmark):
+    def __init__(self, demonstrations_test):
+      self.demonstrations_test = demonstrations_test
+
+    def run(self):
+      with torch.no_grad():
+          states, action_values_desired = self.agent.sample_batch( \
+                  self.demonstrations_test, self.agent.num_eval_episodes)
+          action_values_current = self.agent.online_net(states)
+
+          converted_desired_values = self.agent.convert_values(action_values_desired)
+          converted_current_values = self.agent.convert_values(action_values_current)
+
+          loss = self.agent.calculate_loss(converted_current_values,
+                                           converted_desired_values)
+          scalar_loss = loss.item()
+          result = {f"loss/test": scalar_loss}
+          formatted_result = f"Test-Loss = {scalar_loss}"
+          return result, formatted_result
+
+    def evaluate_loss(self, scalar_loss, converted_current_values, converted_desired_values,
+                    phase="test", logits=False):
+      result = {f"loss/{phase}": scalar_loss}
+      formatted_result = f"Test-Loss = {scalar_loss}"
+
+      return result, formatted_result
+
+    def is_better(self, eval_result1, than_eval_result2):
+        return eval_result1["loss/test"] < than_eval_result2["loss/test"]
+
+class BenchmarkSplitSupervisedLoss(TrainingBenchmark):
+  def __init__(self, demonstrations_test, gradients_and_weights_dir=None):
+    self.demonstrations_test = demonstrations_test
+    self.gradients_and_weights_dir = gradients_and_weights_dir
+
+  def run(self):
+    with torch.no_grad():
+      states, action_values_desired = self.agent.sample_batch( \
+              self.demonstrations_test, self.agent.num_eval_episodes)
+      action_values_current = self.agent.online_net(states)
+
+      converted_desired_values = self.agent.convert_values(action_values_desired)
+      converted_current_values = self.agent.convert_values(action_values_current)
+
+      loss, intermediate_losses = self.agent.calculate_loss(converted_current_values,
+                                                            converted_desired_values,
+                                                            return_intermediate_losses=True)
+
+      return self.evaluate_loss(loss.item(), converted_current_values, converted_desired_values,
+                                intermediate_losses=intermediate_losses)
+
+  def evaluate_loss(self, scalar_loss, converted_current_values, converted_desired_values,
+                    phase="test", logits=False, intermediate_losses=None):
+    result = {f"loss/{phase}": scalar_loss}
+    formatted_result = f"Loss = {scalar_loss}\n          | Mean SE      | Var SE\n"
+
+    if logits:
+      converted_current_values = apply_sigmoid_to_dict(converted_current_values)
+
+    for key in converted_desired_values.keys():
+      pair_wise_diff = converted_desired_values[key] - converted_current_values[key]
+      pair_wise_diff_squared = pair_wise_diff**2
+      error_var, error_mean = torch.var_mean(pair_wise_diff_squared)
+      error_var = error_var.item()
+      error_mean = error_mean.item()
+
+      result[f"squared_error_mean/{key}/{phase}"] = error_mean
+      result[f"squared_error_var/{key}/{phase}"] = error_var
+
+      formatted_result += f"{key:10}| {error_mean:12.8f} | {error_var:12.8f}\n"
+
+    if intermediate_losses is not None:
+      formatted_result += "Intermediate losses:\n"
+      for value_function, intermediate_loss in intermediate_losses.items():
+        result[f"intermediate_loss/{value_function}/{phase}"] = intermediate_loss
+        formatted_result += f"{value_function:10}| {intermediate_loss:12.8f}\n"
+
+    if self.gradients_and_weights_dir:
+      self.save_gradients_and_weights_to_file(
+          self.agent.online_net.named_parameters(), self.agent.steps, phase)
+
+    return result, formatted_result
+
+  def is_better(self, eval_result1, than_eval_result2):
+    return eval_result1["loss/test"] < than_eval_result2["loss/test"]
+
+  def save_gradients_and_weights_to_file(self, named_parameters, steps, phase):
+    weights = dict()
+    gradients = dict()
+
+    for name, param in named_parameters:
+      if "bias" in name:
+        continue
+      weights[name] = param.cpu().detach().numpy()
+      gradients[name] = param.grad.cpu().detach().numpy()
+
+    with open(f"{self.gradients_and_weights_dir}/weights_{phase}_{steps}.npz",
+              "wb") as outfile:
+      np.savez(outfile, **weights)
+    with open(
+        f"{self.gradients_and_weights_dir}/gradients_{phase}_{steps}.npz",
+        "wb") as outfile:
+      np.savez(outfile, **gradients)
+
+
+class ImitationAgent(BaseAgent):
+  def __init__(self, demonstration_collector = None, base_demonstrations_dir=None, demonstration_collector_dir=None, *args, **kwargs):
+    self.demonstration_collector = demonstration_collector
+    self.base_demonstrations_dir = base_demonstrations_dir
+    self.demonstration_collector_dir = demonstration_collector_dir
+    super(ImitationAgent, self).__init__(*args, **kwargs)
+    self.define_training_test_data()
+
+    if self._params["ML"]["ImitationModel"]["SaveGradientsWeights", "", False]:
+      gradients_and_weights_dir = f"{self.agent_save_dir}/gradients_weights"
+      if not os.path.exists(gradients_and_weights_dir):
+        os.mkdir(gradients_and_weights_dir)
+      logging.info("Saving gradients and weights to {}".format(
+          os.path.abspath(gradients_and_weights_dir)))
+    else:
+      gradients_and_weights_dir = None
+
+    self._training_benchmark = BenchmarkSplitSupervisedLoss(
+      self.demonstrations_test, gradients_and_weights_dir=gradients_and_weights_dir)
+    self._training_benchmark.reset(None, self.num_eval_episodes, None, self)
+
+    self.select_loss_function(self._params)
+
+  def define_training_test_data(self):
+    demonstrations = self.demonstration_collector.GetDemonstrationExperiences()
+    random.shuffle(demonstrations)
+    self.demonstrations_train = demonstrations[0:math.floor(len(demonstrations)*self.train_test_ratio)]
+    self.demonstrations_test = demonstrations[math.ceil(len(demonstrations)*self.train_test_ratio):]
+
+  def reset_params(self, params):
+    super(ImitationAgent, self).reset_params(params)
+    self.running_loss_length = params["RunningLossLength", "", 1000]
+    self.num_value_functions = params["NumValueFunctions", "", 3]
+    self.learning_rate = params["LearningRate", "", 0.001]
+    self.train_test_ratio = params["TrainTestRatio", "", 0.8]
+    self.weight_decay = params["WeightDecay", "", 0]
+    self.do_logging = params["DoLogging", "", True]
+
+  def reset_training_variables(self):
+    # Replay memory which is memory-efficient to store stacked frames.
+    self.running_loss = deque(maxlen=self.running_loss_length)
+    self.steps = 0
+    self.best_eval_results = None
+
+  def reset_action_observer(self, env):
+    pass
+
+  @property
+  def observer(self):
+    return self.demonstration_collector.observer
+
+  @property
+  def motion_primitive_behavior(self):
+    return self.demonstration_collector.motion_primitive_behavior
+
+  @property
+  def num_actions(self):
+    return len(self.motion_primitive_behavior.GetMotionPrimitives())
+
+  def init_always(self):
+    super(ImitationAgent, self).init_always()
+    self.init_network()
+
+  def init_network(self):
+    # Target network.
+    self.online_net = Imitation(num_channels=self.observer.observation_space.shape[0],
+                          num_actions=self.num_actions,
+                          num_value_functions=self.num_value_functions,
+                          params=self._params).to(self.device)
+    self.optim = AdamW(
+        self.online_net.parameters(),
+        lr=self.learning_rate,
+        weight_decay=self.weight_decay,
+        betas=(0.9, 0.999), eps=1e-08)
+
+  def clean_pickables(self, pickables):
+    del pickables["_env"]
+    del pickables["_training_benchmark"]
+    del pickables["device"]
+    del pickables["writer"]
+    del pickables["_checkpoint_load"]
+    del pickables["online_net"]
+    del pickables["optim"]
+    del pickables["demonstrations_train"]
+    del pickables["demonstration_collector"]
+
+  def sample_batch(self, demonstrations_list, batch_size):
+    state_size = len(demonstrations_list[0][0])
+    num_action_values = len(demonstrations_list[0][1])
+    states = np.empty((batch_size, state_size), dtype=np.float)
+    action_values = np.empty((batch_size, num_action_values), dtype=np.float)
+
+    indices = np.random.randint(low=0, high=len(demonstrations_list), size=batch_size)
+    for i, index in enumerate(indices):
+      states[i, ...] = demonstrations_list[index][0]
+      action_values[i, ...] = demonstrations_list[index][1]
+
+    states = torch.FloatTensor(states).to(self.device)
+    action_values = torch.FloatTensor(action_values).to(self.device)
+
+    return states, action_values
+
+  def calculate_loss(self, action_values_current, action_values_desired, logits=False,
+                     return_intermediate_losses=False):
+    """
+    If logits=True, the loss must use the sigmoid function before comparing
+    current values to the desired values.
+    Some loss functions (e.g., BCE) need logits for a better numerical
+    stability while training.
+    """
+    # if we have missing actions during demo collections these become nan values
+    # they should not influence the loss and are thus set to current nn output
+    for key in action_values_desired.keys():
+      nans_desired = torch.isnan(action_values_desired[key]).nonzero(as_tuple=False)
+      action_values_desired[key][nans_desired] = action_values_current[key][nans_desired].detach().clone()
+
+    loss = self.selected_loss(action_values_current, action_values_desired, logits,
+                              return_intermediate_losses=return_intermediate_losses)
+    return loss
+
+  def calculate_actions(self, state):
+    with torch.no_grad():
+        state_torch = torch.FloatTensor(state).to(self.device)
+        action_values = self.online_net(state_torch)
+    return action_values.tolist()
+
+  def save_models(self, checkpoint_dir):
+    if not os.path.exists(checkpoint_dir):
+      os.makedirs(checkpoint_dir)
+    torch.save(self.online_net.state_dict(),
+               os.path.join(checkpoint_dir, 'online_net.pth'))
+    online_net_script = torch.jit.script(self.online_net)
+    online_net_script.save(os.path.join(checkpoint_dir, 'online_net_script.pt'))
+
+  def save(self, checkpoint_type="last"):
+    self.demonstration_collector_dir = self.demonstration_collector.GetDirectory()
+    super(ImitationAgent, self).save(checkpoint_type)
+
+  def load_other(self):
+    if not self.demonstration_collector_dir or (
+      not os.path.exists(self.demonstration_collector_dir) and self.base_demonstrations_dir):
+      demonstrations_folder_name = os.path.basename(self.demonstration_collector_dir)
+      self.demonstration_collector_dir = \
+           os.path.abspath(os.path.join(self.base_demonstrations_dir, demonstrations_folder_name))
+      logging.info(f"Resetting demonstration dir to {self.demonstration_collector_dir}")
+    if not os.path.exists(self.demonstration_collector_dir):
+      raise ValueError("Demonstration dir {} not existing".format(self.demonstration_collector_dir))
+    self.demonstration_collector = ActionValuesCollector.load(self.demonstration_collector_dir)
+
+  @property
+  def nn_to_value_converter(self):
+    return self.online_net.nn_to_value_converter
+
+  def load_models(self, checkpoint_dir):
+    try:
+      self.online_net.load_state_dict(
+        torch.load(os.path.join(checkpoint_dir, 'online_net.pth')))
+    except RuntimeError:
+      self.online_net.load_state_dict(
+        torch.load(os.path.join(checkpoint_dir, 'online_net.pth'), map_location=torch.device('cpu')))
+
+  def evaluate_experiences(self, demonstrations):
+    states, action_values_desired = self.sample_batch(demonstrations, len(demonstrations))
+    self.online_net.eval()  # Set to evaluation mode
+    action_values_current = self.online_net(states)
+    self.online_net.train()  # Set back to training mode
+    return  states.detach().numpy(), action_values_desired.detach().numpy(), action_values_current.detach().numpy()
+
+
+  def train_episode(self):
+    states, action_values_desired = self.sample_batch(self.demonstrations_train, self.batch_size)
+
+    self.optim.zero_grad()
+    action_values_current = self.online_net(states)
+
+    converted_desired_values = self.convert_values(action_values_desired)
+    converted_current_values = self.convert_values(action_values_current)
+
+    loss = self.calculate_loss(converted_current_values, converted_desired_values, logits=True)
+    loss.backward()
+    self.optim.step()
+    self.training_log(loss, converted_current_values, converted_desired_values)
+
+  def training_log(self, loss, current_values, desired_values):
+    self.running_loss.append(loss.item())
+    # We log evaluation results along with training frames = 4 * steps.
+    if self.steps % self.summary_log_interval == 0:
+        self.online_net.eval()
+        running_loss_avg = sum(self.running_loss)/len(self.running_loss)
+
+        if self.do_logging:
+          logging.info(f"Training: Loss(i={self.steps}={running_loss_avg})")
+        eval_results, _ = self._training_benchmark.evaluate_loss(
+          running_loss_avg, current_values, desired_values, phase="train", logits=True)
+
+        for eval_result_name, eval_result in eval_results.items():
+          self.writer.add_scalar(eval_result_name, eval_result, self.steps)
+
+        self.online_net.train()
+    if self.steps % self.eval_interval == 0:
+      self.online_net.eval()
+      self.evaluate()
+      self.save("final")
+      self.online_net.train()
+    self.steps += 1
+
+  def select_loss_function(self, params):
+    loss_params = params["ML"]["ImitationModel"]["Loss"]
+    selected_loss = loss_params["SelectedLoss", "", "MeanSquaredErrorLoss"]
+    weights = loss_params["Weights", "", None]
+
+    if selected_loss == "BinaryCrossEntropyLoss":
+      self.selected_loss = LossBCE(weights)
+
+    elif selected_loss == "MeanSquaredErrorLoss":
+      self.selected_loss = LossMSE(weights)
+
+    elif selected_loss == "HuberLoss":
+      delta = loss_params["HuberDelta", "", None]
+      if delta is not None:
+        delta = delta.ConvertToDict()
+      normalize = loss_params["NormalizeHuberLoss", "", False]
+      self.selected_loss = LossHuber(weights, delta, normalize=normalize)
+
+    elif selected_loss == "TukeyLoss":
+      c = loss_params["TukeyLossConstant", "", 0.5]
+      normalize = loss_params["NormalizeTukeyLoss", "", False]
+      self.selected_loss = LossTukey(weights, c=c, normalize=normalize)
+
+    elif selected_loss == "EpsInsensitiveHuberLoss":
+      eps = loss_params["EpsInsensitiveHuberEps", "", 0.01]
+      delta = loss_params["EpsInsensitiveHuberDelta", "", 0.5]
+      normalize = loss_params["NormalizeEpsInsensitiveHuberLoss", "", False]
+      self.selected_loss = LossEpsInsensitiveHuber(weights, eps=eps, delta=delta,
+                                                   normalize=normalize)
+
+    elif selected_loss == "RelativeLoss":
+      self.selected_loss = LossRelative(weights)
+
+    else:
+      logging.warning("Loss not specified or invalid. Using MSE.")
+      self.selected_loss = LossMSE()
+
+  def convert_values(self, raw_values):
+    """
+    Convert a tensor of size (batch size, 3 * number of actions) to a
+    dictionary with an entry for each of the 3 value.
+    """
+    num_actions = raw_values.shape[1] // 3
+    return {
+        "Envelope": raw_values[:, 0:num_actions],
+        "Collision": raw_values[:, num_actions:2*num_actions],
+        "Return": raw_values[:, 2*num_actions:]
+    }
+
+class PolicyImitationAgent(ImitationAgent):
+  def __init__(self, *args, **kwargs):
+    super(PolicyImitationAgent, self).__init__(*args, **kwargs)
+
+    self._training_benchmark = BenchmarkSupervisedLoss(self.demonstrations_test)
+    self._training_benchmark.reset(None, self.num_eval_episodes, None, self)
+
+    self.select_loss_function(self._params)
+
+  def select_loss_function(self, params):
+      loss_params = params["ML"]["ImitationModel"]["Loss"]
+      if loss_params["PolicyCrossEntropyLoss"]:
+        self.selected_loss = LossPolicyCrossEntropy()
+
+      else:
+        self.selected_loss = LossPolicyCrossEntropy()
+        logging.warning("Loss not specified. Using Cross Entropy Policy Loss.")
+
+  def convert_values(self, raw_values):
+    return {
+        "Policy": raw_values
+    }
+
+  def sample_batch(self, demonstrations_list, batch_size):
+    state_size = len(demonstrations_list[0][0])
+    num_actions = len(demonstrations_list[0][2])
+    states = np.empty((batch_size, state_size), dtype=np.float)
+    policy_values = np.empty((batch_size, num_actions), dtype=np.float)
+
+    indices = np.random.randint(low=0, high=len(demonstrations_list), size=batch_size)
+    for i, index in enumerate(indices):
+      states[i, ...] = demonstrations_list[index][0]
+      policy_values[i, ...] = demonstrations_list[index][2]
+
+    states = torch.FloatTensor(states).to(self.device)
+    policy_values = torch.FloatTensor(policy_values).to(self.device)
+
+    return states, policy_values
+
+  def init_network(self):
+    # Target network.
+    self.online_net = PolicyImitation(num_channels=self.observer.observation_space.shape[0],
+                          num_actions=self.num_actions,
+                          params=self._params).to(self.device)
+    self.optim = RMSprop(
+        self.online_net.parameters(),
+        lr=self.learning_rate,
+        weight_decay=self.weight_decay,
+        alpha=0.95,
+        eps=0.00001)
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/iqn_agent.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/iqn_agent.py
index 50e5a935..ebc9643b 100644
--- a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/iqn_agent.py
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/iqn_agent.py
@@ -23,7 +23,7 @@ class IQNAgent(BaseAgent):
   """IQNAgent that can be used in BARK and BARK-ML."""
 
   def __init__(self, env, test_env, params):
-    super(IQNAgent, self).__init__(env, test_env, params)
+    super(IQNAgent, self).__init__(env, params)
 
     self.N = self._params["ML"]["IQNAgent"]["N", "", 64]
     self.N_dash = self._params["ML"]["IQNAgent"]["N_dash", "", 64]
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/loss/BUILD b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/loss/BUILD
new file mode 100644
index 00000000..d0fdcef7
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/loss/BUILD
@@ -0,0 +1,5 @@
+py_library(
+    name = "loss_function",
+    srcs = ["loss_function.py"],
+    visibility = ["//visibility:public"],
+)
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/loss/__init__.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/loss/__init__.py
new file mode 100644
index 00000000..5f85c4f4
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/loss/__init__.py
@@ -0,0 +1,1 @@
+from .loss_function import LossMSE, LossBCE, apply_sigmoid_to_dict
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/loss/loss_function.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/loss/loss_function.py
new file mode 100644
index 00000000..7833e220
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/loss/loss_function.py
@@ -0,0 +1,218 @@
+from torch import sigmoid, nn
+import torch
+
+def apply_sigmoid_to_dict(dict_values):
+  """
+  Returns a new dictionary.
+  """
+  sigmoid_values = dict()
+  for key in dict_values.keys():
+    sigmoid_values[key] = sigmoid(dict_values[key])
+  return sigmoid_values
+
+
+class Loss:
+  """
+  A helper class for calculating loss.
+  """
+  def __init__(self, criterion, weights=None):
+    self._criterion = criterion
+    self._weights = weights
+
+  def __call__(self, current_values, desired_values, logits, return_intermediate_losses=False):
+    """
+    Calculates loss.
+
+    Arguments current_values and desired_values are expected to be
+    dictionaries.
+
+    If return_intermediate_losses==True, return not only the weighted loss,
+    but also a dictionary with a loss for each of the value functions, e.g.:
+      {"Return": 0.03, "Envelope": 0.02, "Collision": 0.01}
+    """
+    if logits:
+      # Transform raw output to values between 0 and 1
+      current_values = apply_sigmoid_to_dict(current_values)
+
+    return self._calculate_weighted_loss(current_values, desired_values, logits,
+                                         return_intermediate_losses=return_intermediate_losses)
+
+  def _select_criterion(self, _logits, _value_func):
+    return self._criterion
+
+  def _calculate_weighted_loss(self, current_values, desired_values, logits,
+                               return_intermediate_losses=False):
+    losses = {}
+    weights_sum = 0
+
+    for value_func in current_values.keys():
+      criterion = self._select_criterion(logits, value_func)
+      weight = self._weights[value_func] if self._weights is not None else 1
+      weights_sum += weight
+
+      loss = criterion(current_values[value_func], desired_values[value_func])
+      losses[value_func] = weight * loss
+
+    weighted_loss = sum(losses.values()) / weights_sum
+
+    if return_intermediate_losses:
+      return weighted_loss, losses
+
+    return weighted_loss
+
+
+class LossMSE(Loss):
+  def __init__(self, weights=None):
+    criterion = nn.MSELoss()
+    super(LossMSE, self).__init__(criterion, weights)
+
+
+class LossBCE(Loss):
+  def __init__(self, weights=None):
+    criterion = nn.BCELoss()
+    super(LossBCE, self).__init__(criterion, weights)
+    self._criterion_logits = nn.BCEWithLogitsLoss()
+
+  def _select_criterion(self, logits, _):
+    if logits:
+      return self._criterion_logits
+    return self._criterion
+
+  def __call__(self, current_values, desired_values, logits, return_intermediate_losses=False):
+    return self._calculate_weighted_loss(current_values, desired_values, logits,
+                                         return_intermediate_losses=return_intermediate_losses)
+
+
+class LossHuber(Loss):
+
+  class HuberCriterion:
+    """
+    torch1.6 does not provide the Huber loss (only available in version 1.9)
+    """
+    def __init__(self, delta, normalize=False):
+      self.delta = delta
+      if normalize:
+        # Make the loss equal to 1 if the absolute difference between target
+        # and prediction is 1
+        max_loss = self._unnormalized_loss(torch.Tensor([1.0]),
+                                           torch.Tensor([0.0])).item()
+        self.normalizing_factor = 1 / max_loss
+      else:
+        self.normalizing_factor = 1
+
+    def _unnormalized_loss(self, current_values, desired_values):
+      error = current_values - desired_values
+      loss = torch.where(error < self.delta,
+                         error**2,
+                         self.delta * (torch.abs(error) - self.delta / 2))
+      return loss.sum() / current_values.data.nelement()
+
+    def __call__(self, current_values, desired_values):
+      return self.normalizing_factor * self._unnormalized_loss(
+          current_values, desired_values)
+
+  def __init__(self, weights=None, delta=None, normalize=False):
+    if delta is None:
+      criterion = self.HuberCriterion(delta=0.5, normalize=normalize)
+      super(LossHuber, self).__init__(criterion, weights)
+    else:
+      self._criterions = {
+          value_func: self.HuberCriterion(delta=d, normalize=normalize)
+          for value_func, d in delta.items()
+      }
+      super(LossHuber, self).__init__(None, weights)
+
+  def _select_criterion(self, _, value_func):
+    if self._criterion is not None:
+      return self._criterion
+    return self._criterions[value_func]
+
+
+class LossTukey(Loss):
+
+  class TukeyCriterion:
+    def __init__(self, c, normalize=False):
+      self.c = c
+      if normalize:
+        # Make the loss equal to 1 if the absolute difference between target
+        # and prediction is 1
+        max_loss = self._unnormalized_loss(torch.Tensor([1.0]),
+                                           torch.Tensor([0.0])).item()
+        self.normalizing_factor = 1 / max_loss
+      else:
+        self.normalizing_factor = 1
+
+    def _unnormalized_loss(self, current_values, desired_values):
+      error = current_values - desired_values
+      const = torch.ones_like(error) * self.c**2 / 6
+      loss = torch.where(error < self.c,
+                         self.c**2 / 6 * (1 - (1 - (error / self.c)**2)**3),
+                         const)
+      return loss.sum() / current_values.data.nelement()
+
+    def __call__(self, current_values, desired_values):
+      return self.normalizing_factor * self._unnormalized_loss(
+          current_values, desired_values)
+
+  def __init__(self, weights=None, c=0.5, normalize=False):
+    criterion = self.TukeyCriterion(c, normalize=normalize)
+    super().__init__(criterion, weights=weights)
+
+
+class LossEpsInsensitiveHuber(Loss):
+
+  class EpsInsensitiveHuberCriterion:
+    def __init__(self, eps, delta, normalize=False):
+      self.eps = eps
+      self.delta = delta
+      if normalize:
+        max_loss = self._unnormalized_loss(torch.Tensor([1.0]),
+                                           torch.Tensor([0.0])).item()
+        self.normalizing_factor = 1 / max_loss
+      else:
+        self.normalizing_factor = 1
+
+    def _unnormalized_loss(self, current_values, desired_values):
+      abs_error = torch.abs(current_values - desired_values)
+      loss = torch.where(
+          abs_error > self.delta,
+          2 * (self.delta - self.eps) * abs_error - self.delta**2 + self.eps**2,
+          torch.Tensor([0.0]))
+      loss = torch.where((abs_error <= self.delta) & (abs_error > self.eps),
+                         (abs_error - self.eps)**2,
+                         loss)
+      return loss.sum() / current_values.data.nelement()
+
+    def __call__(self, current_values, desired_values):
+      return self.normalizing_factor * self._unnormalized_loss(
+          current_values, desired_values)
+
+  def __init__(self, weights=None, eps=0.01, delta=0.5, normalize=False):
+    criterion = self.EpsInsensitiveHuberCriterion(eps, delta, normalize)
+    super().__init__(criterion, weights=weights)
+
+
+class LossRelative(Loss):
+  def __init__(self, weights=None, eps=1e-6):
+    criterion = self._loss
+    self.eps = eps
+    super().__init__(criterion, weights=weights)
+
+  def _loss(self, current_values, desired_values):
+    error = current_values - desired_values
+    loss = torch.abs(error / (desired_values + self.eps))
+    return loss.sum() / current_values.data.nelement()
+
+
+class LossPolicyCrossEntropy(Loss):
+  def __init__(self):
+    pass
+
+  def __call__(self, current_values, desired_values, logits, return_intermediate_losses=False):
+    logsoftmax = nn.LogSoftmax(dim=1)
+    target = desired_values["Policy"]
+    pred = current_values["Policy"]
+    loss = torch.mean(torch.sum(-target*logsoftmax(pred), 1))
+    if return_intermediate_losses:  # Added to match the interface of other losses
+      return loss, None
+    return loss
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/qrdqn_agent.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/qrdqn_agent.py
index a929d111..f8d0fe7c 100644
--- a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/qrdqn_agent.py
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/qrdqn_agent.py
@@ -23,7 +23,7 @@ class QRDQNAgent(BaseAgent):
   """QRDQNAgent that can be used in BARK and BARK-ML."""
 
   def __init__(self, env, test_env, params):
-    super(QRDQNAgent, self).__init__(env, test_env, params)
+    super(QRDQNAgent, self).__init__(env, params)
 
     self.N = self._params["ML"]["QRDQNAgent"]["N", "", 200]
     self.kappa = self._params["ML"]["QRDQNAgent"]["Kappa", "", 1.0]
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/util.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/util.py
new file mode 100644
index 00000000..f0e2c31a
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/util.py
@@ -0,0 +1,16 @@
+# Copyright (c) 2020 Julian Bernhard,
+# Klemens Esterle, Patrick Hart, Tobias Kessler
+#
+# This software is released under the MIT License.
+# https://opensource.org/licenses/MIT
+
+
+def default_training_evaluators():
+  default_config = {"success" : "EvaluatorGoalReached", "collision_other" : "EvaluatorCollisionEgoAgent",
+       "out_of_drivable" : "EvaluatorDrivableArea", "max_steps": "EvaluatorStepCount"}
+  return default_config
+
+def default_terminal_criteria(max_episode_steps):
+  terminal_when = {"collision_other" : lambda x: x, "out_of_drivable" : lambda x: x, \
+        "max_steps": lambda x : x>max_episode_steps, "success" : lambda x: x}
+  return terminal_when
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/BUILD b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/BUILD
index aee92c23..c736bccf 100644
--- a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/BUILD
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/BUILD
@@ -6,6 +6,7 @@ py_library(
         ":fqf",
         ":iqn",
         ":qrdqn",
+        ":imitation"
     ],
     visibility = ["//visibility:public"],
 )
@@ -25,6 +26,11 @@ py_library(
     ],
 )
 
+py_library(
+    name = "imitation",
+    srcs = ["imitation.py"]
+)
+
 py_library(
     name = "iqn",
     srcs = ["iqn.py"],
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/__init__.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/__init__.py
index b468b8df..ea1725f5 100644
--- a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/__init__.py
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/__init__.py
@@ -1,3 +1,4 @@
-from .fqf import FQF  # pylint: disable=unused-import
-from .iqn import IQN  # pylint: disable=unused-import
-from .qrdqn import QRDQN  # pylint: disable=unused-import
+from .fqf import FQF
+from .iqn import IQN
+from .qrdqn import QRDQN
+from .imitation import Imitation, PolicyImitation
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/imitation.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/imitation.py
new file mode 100644
index 00000000..45e31782
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/imitation.py
@@ -0,0 +1,90 @@
+import torch
+from torch import nn, sigmoid, softmax
+from collections import OrderedDict
+from bark_ml.core.value_converters import *
+
+def init_weights(m):
+    if type(m) == nn.Linear:
+        nn.init.xavier_uniform_(m.weight)
+        m.bias.data.fill_(0.01)
+
+class Imitation(nn.Module):
+  def __init__(self,
+               num_channels,
+               num_actions,
+               num_value_functions,
+               params):
+    super(Imitation, self).__init__()
+    self.num_channels = num_channels
+    self.num_actions = num_actions
+    self.num_value_functions = num_value_functions
+    self.value_converter = NNToValueConverterSequential(self.num_actions)
+    self.layer_dims = params["ML"]["ImitationModel"]["EmbeddingDims", "", [256, 256, 256]]
+    self.droput_p = params["ML"]["ImitationModel"]["DropoutProbability", "", 0]
+
+    self.net = nn.Sequential(self.make_ordered_layer_dict(self.layer_dims))
+    self.net.apply(init_weights)
+
+  def make_ordered_layer_dict(self, layer_dims):
+      tuple_list = []
+      last_dim = self.num_channels
+      for idx, layer in enumerate(layer_dims):
+          current_dim = layer
+          tuple_list.append((f"layer{idx}", nn.Linear(last_dim, current_dim)))
+          tuple_list.append((f"relu{idx}", nn.ReLU()))
+          if self.droput_p != 0:
+            tuple_list.append((f"dropout{idx}", nn.Dropout(p=self.droput_p)))
+          last_dim = current_dim
+      tuple_list.append(("output", nn.Linear(last_dim, self.num_actions*self.num_value_functions)))
+      return OrderedDict(tuple_list)
+
+  @torch.jit.unused
+  @property
+  def nn_to_value_converter(self):
+    return self.value_converter
+
+  def forward(self, states):
+    action_values = self.net(states)
+    if not self.training:
+      # Evaluation phase, output values between 0 and 1
+      action_values = sigmoid(action_values)
+    return action_values
+
+class PolicyImitation(nn.Module):
+  def __init__(self,
+               num_channels,
+               num_actions,
+               params):
+    super(PolicyImitation, self).__init__()
+    self.num_channels = num_channels
+    self.num_actions = num_actions
+    self.value_converter = NNToValueConverterPolicy(self.num_actions)
+    self.layer_dims = params["ML"]["ImitationModel"]["EmbeddingDims", "", [256, 256, 256]]
+    self.droput_p = params["ML"]["ImitationModel"]["DropoutProbability", "", 0]
+
+    self.net = nn.Sequential(self.make_ordered_layer_dict(self.layer_dims))
+    self.net.apply(init_weights)
+
+  def make_ordered_layer_dict(self, layer_dims):
+      tuple_list = []
+      last_dim = self.num_channels
+      for idx, layer in enumerate(layer_dims):
+          current_dim = layer
+          tuple_list.append((f"layer{idx}", nn.Linear(last_dim, current_dim)))
+          tuple_list.append((f"relu{idx}", nn.ReLU()))
+          if self.droput_p != 0:
+            tuple_list.append((f"dropout{idx}", nn.Dropout(p=self.droput_p)))
+          last_dim = current_dim
+      tuple_list.append(("last_linear", nn.Linear(last_dim, self.num_actions)))
+      return OrderedDict(tuple_list)
+
+  @property
+  def nn_to_value_converter(self):
+    return self.value_converter
+
+  def forward(self, states):
+    action_values = self.net(states)
+    if not self.training:
+      # Evaluation phase, output values between 0 and 1
+      action_values = softmax(action_values, dim=1)
+    return action_values
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/BUILD b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/BUILD
new file mode 100644
index 00000000..7cb0531b
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/BUILD
@@ -0,0 +1,10 @@
+cc_library(
+  name = "nn_to_value_converter",
+  hdrs = [
+    "nn_to_value_converter.hpp",
+    "nn_to_value_converter_sequential.hpp",
+    "nn_to_value_converter_policy.hpp"
+  ],
+  linkstatic=False,
+  visibility = ["//visibility:public"],
+)
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter.hpp b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter.hpp
new file mode 100644
index 00000000..b4c00482
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter.hpp
@@ -0,0 +1,60 @@
+
+
+
+// Copyright (c) fortiss GmbH
+//
+// Authors: Patrick Hart, Julian Bernhard, Klemens Esterle
+// Tobias Kessler and Mansoor Nasir
+//
+// This work is licensed under the terms of the MIT license.
+// For a copy, see <https://opensource.org/licenses/MIT>.
+
+
+#ifndef BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_NN_TO_VALUE_CONVERTER_NN_TO_VALUE_CONVERTER_HPP_
+#define BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_NN_TO_VALUE_CONVERTER_NN_TO_VALUE_CONVERTER_HPP_
+
+#include <unordered_map>
+#include <vector>
+#include <memory>
+
+#ifdef GOOGLE_STRIP_LOG
+#define LOG_INFO VLOG(4)
+#define LOG_ERROR LOG(ERROR)
+#define LOG_FATAL LOG(FATAL)
+#else
+#define LOG 
+#define LOG_INFO std::cout
+#define LOG_ERROR std::cerr
+#define LOG_FATAL std::cerr
+#endif
+
+
+namespace bark_ml {
+namespace lib_fqf_iqn_qrdqn {
+
+typedef enum ValueType {
+    Return = 0,
+    EnvelopeRisk = 1,
+    CollisionRisk = 2,
+    Policy = 3
+} ValueType;
+
+class NNToValueConverter {
+ public:
+  NNToValueConverter(const unsigned& num_actions) : num_actions_(num_actions) {}
+
+  virtual ~NNToValueConverter() {}
+
+  virtual std::unordered_map<ValueType, std::vector<double>> ConvertToValueMap(const std::vector<float>& nn_output) const = 0;
+
+  unsigned GetNumActions() const { return num_actions_; }
+protected:
+ unsigned num_actions_;
+};
+
+typedef std::shared_ptr<NNToValueConverter> NNToValueConverterPtr;
+
+}  // namespace lib_fqf_iqn_qrdqn
+}  // namespace bark_ml
+
+#endif //BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_NN_TO_VALUE_CONVERTER_NN_TO_VALUE_CONVERTER_HPP_
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter_policy.hpp b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter_policy.hpp
new file mode 100644
index 00000000..952abd44
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter_policy.hpp
@@ -0,0 +1,46 @@
+
+
+
+// Copyright (c) fortiss GmbH
+//
+// Authors: Patrick Hart, Julian Bernhard, Klemens Esterle
+// Tobias Kessler and Mansoor Nasir
+//
+// This work is licensed under the terms of the MIT license.
+// For a copy, see <https://opensource.org/licenses/MIT>.
+
+
+#ifndef BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_NN_TO_VALUE_CONVERTER_NN_TO_VALUE_CONVERTER_POLICY_HPP_
+#define BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_NN_TO_VALUE_CONVERTER_NN_TO_VALUE_CONVERTER_POLICY_HPP_
+
+#include <math.h>
+#include "bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter.hpp"
+
+
+namespace bark_ml {
+namespace lib_fqf_iqn_qrdqn {
+
+class NNToValueConverterPolicy : public NNToValueConverter {
+ public:
+  NNToValueConverterPolicy(const unsigned& num_actions) : NNToValueConverter(num_actions) {}
+
+  virtual ~NNToValueConverterPolicy() {}
+
+  virtual std::unordered_map<ValueType, std::vector<double>> ConvertToValueMap(const std::vector<float>&  nn_output) const {
+    std::unordered_map<ValueType, std::vector<double>> value_map;
+
+    if(nn_output.size() != num_actions_) {
+      LOG_FATAL << "Invalid neural network output with size of " << nn_output.size() << " for num_actions = " << num_actions_;
+    }
+    // Only a policy output for this conversion type
+    value_map[ValueType::Policy] = std::vector<double>(nn_output.begin(), nn_output.begin()+num_actions_);
+    
+    return value_map;
+  };
+
+};
+
+}  // namespace lib_fqf_iqn_qrdqn
+}  // namespace bark_ml
+
+#endif //BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_NN_TO_VALUE_CONVERTER_NN_TO_VALUE_CONVERTER_POLICY_HPP_
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter_sequential.hpp b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter_sequential.hpp
new file mode 100644
index 00000000..e21f474c
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter_sequential.hpp
@@ -0,0 +1,49 @@
+
+
+
+// Copyright (c) fortiss GmbH
+//
+// Authors: Patrick Hart, Julian Bernhard, Klemens Esterle
+// Tobias Kessler and Mansoor Nasir
+//
+// This work is licensed under the terms of the MIT license.
+// For a copy, see <https://opensource.org/licenses/MIT>.
+
+
+#ifndef BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_NN_TO_VALUE_CONVERTER_NN_TO_VALUE_CONVERTER_SEQUENTIAL_HPP_
+#define BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_NN_TO_VALUE_CONVERTER_NN_TO_VALUE_CONVERTER_SEQUENTIAL_HPP_
+
+#include "bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter.hpp"
+
+
+namespace bark_ml {
+namespace lib_fqf_iqn_qrdqn {
+
+class NNToValueConverterSequential : public NNToValueConverter {
+ public:
+  NNToValueConverterSequential(const unsigned& num_actions) : NNToValueConverter(num_actions) {}
+
+  virtual ~NNToValueConverterSequential() {}
+
+  virtual std::unordered_map<ValueType, std::vector<double>> ConvertToValueMap(const std::vector<float>&  nn_output) const {
+    std::unordered_map<ValueType, std::vector<double>> value_map;
+
+    if(nn_output.size() != num_actions_*3) {
+      LOG_FATAL << "Invalid neural network output with size of " << nn_output.size() << " for num_actions = " << num_actions_;
+    }
+    // Convention that first num_actions elements are return 
+    value_map[ValueType::EnvelopeRisk] = std::vector<double>(nn_output.begin(), nn_output.begin()+num_actions_);
+    //... second num action elements are envelope risk
+    value_map[ValueType::CollisionRisk] = std::vector<double>(nn_output.begin()+num_actions_, nn_output.begin()+2*num_actions_);
+    // ... and third num action elements are collision risk
+    value_map[ValueType::Return] = std::vector<double>(nn_output.begin()+2*num_actions_, nn_output.begin()+3*num_actions_);
+
+    return value_map;
+  };
+
+};
+
+}  // namespace lib_fqf_iqn_qrdqn
+}  // namespace bark_ml
+
+#endif //BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_NN_TO_VALUE_CONVERTER_NN_TO_VALUE_CONVERTER_SEQUENTIAL_HPP_
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_wrapper/BUILD b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_loader/BUILD
similarity index 71%
rename from bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_wrapper/BUILD
rename to bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_loader/BUILD
index f690d370..3c9fdaad 100644
--- a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_wrapper/BUILD
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_loader/BUILD
@@ -9,7 +9,7 @@ pybind_extension(
   copts = ["-Iexternal/torchcpp/torch/include",
            "-Lexternal/torchcpp/torch/lib",
            "-D_GLIBCXX_USE_CXX11_ABI=0"],
-   linkopts = ["-ltorch_cpu","-lc10"],
+   linkopts = ["-ltorch_cpu","-lc10", "-ltorch"],
    visibility = ["//visibility:public"],
 )
 
@@ -17,10 +17,6 @@ cc_library(
   name = "model_loader",
   hdrs = ["model_loader.hpp"],
   deps = ["@torchcpp//:lib"],
-  copts = ["-Iexternal/torchcpp/torch/include",
-           "-Lexternal/torchcpp/torch/lib",
-           "-D_GLIBCXX_USE_CXX11_ABI=0"],
-
-  linkopts = ["-ltorch_cpu","-lc10"],
   linkstatic=False,
+  visibility = ["//visibility:public"],
 )
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_loader/model_loader.hpp b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_loader/model_loader.hpp
new file mode 100644
index 00000000..2b709c24
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_loader/model_loader.hpp
@@ -0,0 +1,99 @@
+// Copyright (c) fortiss GmbH
+//
+// Authors: Patrick Hart, Julian Bernhard, Klemens Esterle
+// Tobias Kessler and Mansoor Nasir
+//
+// This work is licensed under the terms of the MIT license.
+// For a copy, see <https://opensource.org/licenses/MIT>.
+
+
+#ifndef BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_LOADER_HPP_
+#define BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_LOADER_HPP_
+
+#include <torch/script.h>
+#include <ATen/Parallel.h>
+
+#include <chrono>
+#include <iostream>
+#include <memory>
+// Logging based on availability of glog
+#ifdef GOOGLE_STRIP_LOG
+#define LOG_INFO VLOG(4)
+#define LOG_ERROR LOG(ERROR)
+#else
+#define LOG 
+#define LOG_INFO std::cout
+#define LOG_ERROR std::cerr
+#endif
+
+namespace bark_ml {
+namespace lib_fqf_iqn_qrdqn {
+
+// Load a pytorch trained network (nn.Module) saved from
+// python and perform an inference
+class ModelLoader {
+ public:
+  ModelLoader() 
+    : module_loaded_(false) {
+  }
+
+  bool LoadModel(const std::string& model_filename) {
+    if (module_loaded_) {
+      return true; 
+    }
+    LOG_INFO << "Trying to load model from file: " << model_filename << "\n";
+    try {
+      module_ = torch::jit::load(model_filename);
+      module_.eval();
+      module_loaded_ = true;
+    }
+    catch (const c10::Error& e) {
+      LOG_ERROR << "Error loading the model: " << e.msg();
+    }
+    at::init_num_threads();
+    LOG_INFO << "ML: get_num_threads(): " << at::get_num_threads() << ", get_num_interop_threads(): " << at::get_num_interop_threads();
+    return module_loaded_;
+  }
+
+  std::vector<float> Inference(std::vector<float> state) {
+    if (!module_loaded_)
+      throw std::runtime_error("Model not loaded!");
+
+    std::vector<torch::jit::IValue> inputs;
+    inputs.push_back(torch::tensor(state).view({1, state.size()}));
+
+    at::Tensor torch_output;
+    try {
+      // run the inference
+      torch_output = module_.forward(inputs).toTensor();
+
+    } catch (const c10::Error& e) {
+      std::cerr << e.msg();
+      throw std::runtime_error(e.msg());
+    }
+
+    //copy the data from torch tensor to std vector
+    std::vector<float> output(torch_output.data_ptr<float>(), torch_output.data_ptr<float>() + torch_output.numel());
+
+    return output;
+  }
+
+  double TimeMeasuredInference(std::vector<float> state) {
+    auto start = std::chrono::high_resolution_clock::now();
+    const auto result = Inference(state);
+    auto finish = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<double, std::micro> elapsed = finish - start;
+    double last_inference_time = elapsed.count();
+    return last_inference_time;
+  }
+
+ private:
+  bool module_loaded_;
+  torch::jit::script::Module module_;
+
+};
+
+}  // namespace lib_fqf_iqn_qrdqn
+}  // namespace bark_ml
+
+#endif //BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_LOADER_HPP_
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_wrapper/pytorch_script_wrapper.cpp b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_loader/pytorch_script_wrapper.cpp
similarity index 66%
rename from bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_wrapper/pytorch_script_wrapper.cpp
rename to bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_loader/pytorch_script_wrapper.cpp
index 4205c7c5..bcd9a89a 100644
--- a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_wrapper/pytorch_script_wrapper.cpp
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_loader/pytorch_script_wrapper.cpp
@@ -1,4 +1,4 @@
-#include "bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_wrapper/model_loader.hpp"
+#include "bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_loader/model_loader.hpp"
 #include <iostream>
 #include <memory>
 
@@ -11,7 +11,8 @@ using bark_ml::lib_fqf_iqn_qrdqn::ModelLoader;
 
 PYBIND11_MODULE(pytorch_script_wrapper, m) {
     py::class_<ModelLoader>(m, "ModelLoader")
-        .def(py::init<const std::string &, const long, const long>())
+        .def(py::init<>())
         .def("Inference", &ModelLoader::Inference, "Perform the inference for a given state")
-        .def("LoadModel", &ModelLoader::LoadModel, "Loads the torch cpp script model");
+        .def("LoadModel", &ModelLoader::LoadModel, "Loads the torch cpp script model")
+        .def("TimeMeasuredInference", &ModelLoader::TimeMeasuredInference);
 }
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_wrapper/model_loader.hpp b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_wrapper/model_loader.hpp
deleted file mode 100644
index a40f2fc1..00000000
--- a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_wrapper/model_loader.hpp
+++ /dev/null
@@ -1,86 +0,0 @@
-// Copyright (c) fortiss GmbH
-//
-// Authors: Patrick Hart, Julian Bernhard, Klemens Esterle
-// Tobias Kessler and Mansoor Nasir
-//
-// This work is licensed under the terms of the MIT license.
-// For a copy, see <https://opensource.org/licenses/MIT>.
-
-
-#ifndef BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_LOADER_HPP_
-#define BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_LOADER_HPP_
-
-#include <torch/script.h>
-
-#include <iostream>
-#include <memory>
-
-namespace bark_ml {
-namespace lib_fqf_iqn_qrdqn {
-
-// Load a pytorch trained network (nn.Module) saved from
-// python and perform an inference
-class ModelLoader {
- public:
-  ModelLoader(const std::string& model_filename, const long action_space, const long state_space) 
-    : model_filename_(model_filename), module_loaded_(false), action_space_(action_space),
-     state_space_(state_space) {
-  }
-
-  bool LoadModel() {
-    if (module_loaded_) {
-      return true; // module already loaded
-    }
-
-    // load the torch script module saved from python
-    try {
-      module_ = torch::jit::load(model_filename_);
-      module_loaded_ = true;
-    }
-    catch (const c10::Error& e) {
-      std::cerr << "Error loading the model\n";
-    }
-
-    return module_loaded_;
-  }
-
-  std::vector<double> Inference(std::vector<double> state) {
-    std::vector<double> actions(action_space_);
-
-    if (!module_loaded_)
-      throw std::runtime_error("Model not loaded!");
-
-    std::vector<torch::jit::IValue> inputs;
-    inputs.push_back(torch::tensor(state).view({1, state_space_}));
-
-    at::Tensor predicted_actions;
-    try {
-      // run the inference
-      predicted_actions = module_.forward(inputs).toTensor();
-
-    } catch (const c10::Error& e) {
-      std::cerr << e.msg();
-      throw std::runtime_error(e.msg());
-    }
-
-    //copy the data from torch tensor to std vector
-    auto options = torch::TensorOptions().dtype(torch::kFloat32);
-    torch::Tensor tharray = torch::from_blob(actions.data(), {1, action_space_}, options);
-    tharray.copy_(predicted_actions);
-
-    return actions;
-  }
-
- private:
-  std::string model_filename_;
-  long action_space_;
-  long state_space_;
-  bool module_loaded_;
-  torch::jit::script::Module module_;
-
-};
-
-}  // namespace lib_fqf_iqn_qrdqn
-}  // namespace bark_ml
-
-#endif //BARK_ML_LIB_FQF_IQN_QRDQN_MODEL_LOADER_HPP_
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_tf_agents/agents/tfa_agent.py b/bark_ml/library_wrappers/lib_tf_agents/agents/tfa_agent.py
index 5f725d13..348a73d1 100644
--- a/bark_ml/library_wrappers/lib_tf_agents/agents/tfa_agent.py
+++ b/bark_ml/library_wrappers/lib_tf_agents/agents/tfa_agent.py
@@ -114,7 +114,7 @@ def Plan(self, dt, observed_world):
       # NOTE: we need to store the observer differently
       observed_state = self._environment._observer.Observe(
         observed_world)
-      self._action = self.Act(observed_state)
+      self._action = self.Act(np.float32(observed_state))
     # NOTE: BARK expects (m, 1) actions
     action = self._action
     if isinstance(self.action_space, BoundedContinuous):
diff --git a/bark_ml/library_wrappers/lib_tf_agents/runners/tfa_runner.py b/bark_ml/library_wrappers/lib_tf_agents/runners/tfa_runner.py
index d1a8d4ce..341ae7a9 100644
--- a/bark_ml/library_wrappers/lib_tf_agents/runners/tfa_runner.py
+++ b/bark_ml/library_wrappers/lib_tf_agents/runners/tfa_runner.py
@@ -127,7 +127,7 @@ def RunEpisode(self, render=True):
     reward = 0
     while not is_terminal:
       action_step = self._agent._eval_policy.action(
-        ts.transition(state, reward=0.0, discount=1.0))
+        ts.transition(np.float32(state), reward=0.0, discount=1.0))
       action = self.ReshapeActionIfRequired(action_step)
 
       state, reward, is_terminal, info = self._environment.step(action)
diff --git a/bark_ml/observers/BUILD b/bark_ml/observers/BUILD
index 26eea3d2..114db9cf 100644
--- a/bark_ml/observers/BUILD
+++ b/bark_ml/observers/BUILD
@@ -26,3 +26,16 @@ cc_library(
   ],
   visibility = ["//visibility:public"],
 )
+
+cc_library(
+  name = "include",
+  hdrs = [
+    "base_observer.hpp",
+    "nearest_observer.hpp",
+  ],
+  deps=[
+    "//bark_ml/commons:spaces",
+    "//bark_ml/commons:commons"
+  ],
+  visibility = ["//visibility:public"],
+)
diff --git a/bark_ml/observers/base_observer.hpp b/bark_ml/observers/base_observer.hpp
index 2d20c577..bafa76c0 100644
--- a/bark_ml/observers/base_observer.hpp
+++ b/bark_ml/observers/base_observer.hpp
@@ -31,9 +31,8 @@ namespace observers {
 using bark::commons::ParamsPtr;
 using spaces::Box;
 using bark::world::WorldPtr;
-using bark::world::ObservedWorldPtr;
-using ObservedState = Eigen::Matrix<float, 1, Eigen::Dynamic>;
-using State = Eigen::Matrix<double, Eigen::Dynamic, 1>;
+using bark::world::ObservedWorld;
+using ObservedState = Eigen::Matrix<double, 1, Eigen::Dynamic>;
 
 /**
  * @brief  Base class for the Observer.
@@ -44,14 +43,17 @@ class BaseObserver {
     params_(params) {}
 
   virtual ObservedState Observe(
-    const ObservedWorldPtr& observed_world) const = 0;
+    const ObservedWorld& observed_world) const = 0;
   virtual WorldPtr Reset(const WorldPtr& world) = 0;
   virtual Box<double> ObservationSpace() const = 0;
+  ParamsPtr GetParams() const { return params_; }
 
  private:
   ParamsPtr params_;
 };
 
+typedef std::shared_ptr<bark_ml::observers::BaseObserver> ObserverPtr;
+
 }  // namespace observers
 }  // namespace bark_ml
 
diff --git a/bark_ml/observers/nearest_observer.hpp b/bark_ml/observers/nearest_observer.hpp
index f8ded420..9da63f2f 100644
--- a/bark_ml/observers/nearest_observer.hpp
+++ b/bark_ml/observers/nearest_observer.hpp
@@ -34,15 +34,26 @@ using commons::Norm;
 using spaces::Matrix_t;
 using bark::world::AgentMap;
 using bark::world::AgentPtr;
+using bark::world::objects::AgentId;
 using bark::world::WorldPtr;
+using bark::world::goal_definition::GoalDefinitionStateLimitsFrenet;
+using bark::world::map::LaneCorridorPtr;
 using bark::world::ObservedWorldPtr;
 using bark::geometry::Point2d;
+using bark::commons::transformation::FrenetState;
+using bark::commons::transformation::FrenetStateDifference;
+using bark::geometry::B_2PI;
+using bark::geometry::B_PI;
+using bark::geometry::Line;
 using bark::geometry::Distance;
-using bark::geometry::Norm0To2PI;
+using bark::geometry::NormToPI;
 using bark::models::dynamic::StateDefinition::X_POSITION;
 using bark::models::dynamic::StateDefinition::Y_POSITION;
 using bark::models::dynamic::StateDefinition::THETA_POSITION;
 using bark::models::dynamic::StateDefinition::VEL_POSITION;
+using ObservedState = Eigen::Matrix<double, 1, Eigen::Dynamic>;
+using bark::commons::transformation::FrenetPosition;
+using State = Eigen::Matrix<double, Eigen::Dynamic, 1>;
 
 
 /**
@@ -51,79 +62,108 @@ using bark::models::dynamic::StateDefinition::VEL_POSITION;
 class NearestObserver : public BaseObserver {
  public:
   explicit NearestObserver(const ParamsPtr& params) :
-    BaseObserver(params),
-    min_x_(0.), max_x_(100.),
-    min_y_(0.), max_y_(100.),
-    min_theta_(0.), max_theta_(2*3.14) {
+    BaseObserver(params) {
       nearest_agent_num_ =
         params->GetInt(
           "ML::NearestObserver::NNearestAgents", "Nearest agents number", 4);
-      min_vel_ = params->GetReal("ML::NearestObserver::MinVel", "", 0.0);
-      max_vel_ = params->GetReal("ML::NearestObserver::MaxVel", "", 50.0);
+      min_vel_lon_ = params->GetReal("ML::NearestObserver::MinVelLon", "", -30.0);
+      max_vel_lon_ = params->GetReal("ML::NearestObserver::MaxVelLon", "", 30.0);
+      min_vel_lat_ = params->GetReal("ML::NearestObserver::MinVelLat", "", -10.0);
+      max_vel_lat_ = params->GetReal("ML::NearestObserver::MaxVelLat", "", 10.0);
       max_dist_ = params->GetReal("ML::NearestObserver::MaxDist", "", 75.0);
-      state_size_ = params->GetInt("ML::NearestObserver::StateSize", "", 4);
-      observation_len_ = nearest_agent_num_ * state_size_;
+      min_s_ = params->GetReal("ML::NearestObserver::MinS", "", -100.0);
+      min_abs_s_ = params->GetReal("ML::NearestObserver::MinAbsS", "", 0.0);
+      max_s_ = params->GetReal("ML::NearestObserver::MaxS", "", 100.0);
+      min_d_ = params->GetReal("ML::NearestObserver::MinD", "", -20.0);
+      max_d_ = params->GetReal("ML::NearestObserver::MaxD", "", 20.0);
+      min_theta_ = params->GetReal("ML::NearestObserver::MinTheta", "", -B_PI);
+      max_theta_ = params->GetReal("ML::NearestObserver::MaxTheta", "", B_PI);
+      // Ego agent has s,d, vlat, vlon, theta, laneid other agents have del s, del d, del vlat, del vlon
+      observation_len_ = nearest_agent_num_ * 4 + 6; 
   }
 
   double Norm(const double val, const double mi, const double ma) const {
-    return (val - mi)/(ma - mi);
+    LOG_IF_EVERY_N(WARNING, (val < mi), 100) << "Val=" << val << " < Lower Limit=" << mi;
+    LOG_IF_EVERY_N(WARNING, (val > ma), 100) << "Val=" << val << " > Upper Limit=" << ma;
+    // Normalize to be within -1 to 1
+    return (val - (mi+ma)/2.0)/((ma-mi)/2.0);
   }
 
-  ObservedState FilterState(const State& state) const {
-    ObservedState ret_state(1, state_size_);
-    const double normalized_angle = Norm0To2PI(state(THETA_POSITION));
-    ret_state << Norm(state(X_POSITION), min_x_, max_x_),
-                 Norm(state(Y_POSITION), min_y_, max_y_),
+  ObservedState GetEgoState(const ObservedWorld& observed_world) const {
+    const auto ego_state = observed_world.CurrentEgoState();
+    const auto ego_pos = observed_world.CurrentEgoPosition();
+    const auto corridor_and_idx = observed_world.GetEgoAgent()
+      ->GetRoadCorridor()->GetNearestLaneCorridorAndIndex(ego_pos);
+    const auto& ego_corridor = corridor_and_idx.first;
+    const unsigned corr_idx = corridor_and_idx.second;
+    FrenetState current_ego_frenet(ego_state, ego_corridor->GetCenterLine());
+    ObservedState ego_nn_state(1, 6);
+    const double normalized_angle = NormToPI(current_ego_frenet.angle);
+    LOG_IF_EVERY_N(WARNING, (current_ego_frenet.lon < min_abs_s_), 100)
+         << "Sego=" << current_ego_frenet.lon << " < Min abs s=" << min_abs_s_;
+    ego_nn_state << Norm(current_ego_frenet.lon - min_abs_s_, min_s_, max_s_),
+                 Norm(current_ego_frenet.lat, min_d_, max_d_),
                  Norm(normalized_angle, min_theta_, max_theta_),
-                 Norm(state(VEL_POSITION), min_vel_, max_vel_);
-    return ret_state;
+                 Norm(current_ego_frenet.vlon, min_vel_lon_, max_vel_lon_),
+                 Norm(current_ego_frenet.vlat, min_vel_lat_, max_vel_lat_),
+                corr_idx;
+    return ego_nn_state;
   }
 
-  ObservedState Observe(const ObservedWorldPtr& observed_world) const {
-    int row_idx = 0;
+  ObservedState GetOtherAgentState(const AgentId& agent_id, const ObservedWorld& observed_world) const {
+    const auto ego_pos = observed_world.CurrentEgoPosition();
+    const auto corridor_and_idx = observed_world.GetEgoAgent()
+            ->GetRoadCorridor()->GetNearestLaneCorridorAndIndex(ego_pos);
+
+    const auto ego_state = observed_world.CurrentEgoState();
+    FrenetState current_ego_frenet(ego_state, corridor_and_idx.first->GetCenterLine());
+    const auto ego_shape = observed_world.GetEgoAgent()->GetShape();
+    const auto& other_state = observed_world.GetAgent(agent_id)->GetCurrentState();
+    FrenetState other_frenet(other_state, corridor_and_idx.first->GetCenterLine());
+    const auto other_shape = observed_world.GetAgent(agent_id)->GetShape();
+    FrenetStateDifference state_difference(current_ego_frenet, ego_shape, other_frenet, other_shape);
+
+    ObservedState other_nn_state(1, 4);
+    other_nn_state << Norm(state_difference.lon_zeroed ? 0.0 : state_difference.lon, min_s_, max_s_),
+                 Norm(state_difference.lat_zeroed ? 0.0 : state_difference.lat, min_d_, max_d_),
+                 Norm(state_difference.vlon, min_vel_lon_, max_vel_lon_),
+                 Norm(state_difference.vlat, min_vel_lat_, max_vel_lat_);
+    return other_nn_state;
+  }
+
+  ObservedState Observe(const ObservedWorld& observed_world) const {
+    // find near agents (n)
+    AgentMap nearest_agents = observed_world.GetNearestAgents(
+      observed_world.CurrentEgoPosition(), nearest_agent_num_ + 1);
+
+    // Build state
+    int state_start_idx = 0;
     ObservedState state(1, observation_len_);
     state.setZero();
 
-    // find near agents (n)
-    AgentMap nearest_agents = observed_world->GetNearestAgents(
-      observed_world->CurrentEgoPosition(), nearest_agent_num_);
+    // add ego agent state
+    ObservedState obs_ego_agent_state = GetEgoState(observed_world);
+    state.block(0, state_start_idx, 1, obs_ego_agent_state.cols()) = obs_ego_agent_state;
+    state_start_idx = obs_ego_agent_state.cols();
 
-    // sort agents by distance and distance < max_dist_
-    std::map<double, AgentPtr, std::greater<double>> distance_agent_map;
+    // add other states
     for (const auto& agent : nearest_agents) {
+      if(agent.first == observed_world.GetEgoAgentId()) continue;
+
       const auto& agent_state = agent.second->GetCurrentPosition();
       double distance = Distance(
-        observed_world->CurrentEgoPosition(), agent_state);
-      if (distance < max_dist_)
-        distance_agent_map[distance] = agent.second;
-    }
+        observed_world.CurrentEgoPosition(), agent_state);
+      if (distance > max_dist_) continue;
 
-    // add ego agent state
-    ObservedState obs_ego_agent_state =
-      FilterState(observed_world->CurrentEgoState());
-    state.block(0, row_idx*state_size_, 1, state_size_) = obs_ego_agent_state;
-    row_idx++;
-
-    // add other states
-    for (const auto& agent : distance_agent_map) {
-      if (agent.second->GetAgentId() != observed_world->GetEgoAgentId()) {
-        ObservedState other_agent_state =
-          FilterState(agent.second->GetCurrentState());
-        state.block(0, row_idx*state_size_, 1, state_size_) = other_agent_state;  // # pylint: disable=unused-import
-        row_idx++;
-      }
+      ObservedState other_agent_state = GetOtherAgentState(agent.second->GetAgentId(),
+                      observed_world);
+      state.block(0, state_start_idx, 1, other_agent_state.cols()) = other_agent_state;  // NOLINT
+      state_start_idx += other_agent_state.cols();
     }
     return state;
   }
 
   WorldPtr Reset(const WorldPtr& world) {
-    const auto& x_y = world->BoundingBox();
-    Point2d bb0 = x_y.first;
-    Point2d bb1 = x_y.second;
-    min_x_ = bb0.get<0>();
-    max_x_ = bb1.get<0>();
-    min_y_ = bb0.get<1>();
-    max_y_ = bb1.get<1>();
     return world;
   }
 
@@ -137,9 +177,11 @@ class NearestObserver : public BaseObserver {
   }
 
  private:
-  int state_size_, nearest_agent_num_, observation_len_;
-  double min_theta_, max_theta_, min_vel_, max_vel_, max_dist_,
-         min_x_, max_x_, min_y_, max_y_;
+  int nearest_agent_num_, observation_len_;
+  double min_theta_, max_theta_, min_vel_lon_,
+         max_vel_lon_, min_vel_lat_,
+         max_vel_lat_, max_dist_,
+         min_d_, max_d_, min_s_, max_s_, min_abs_s_;
 };
 
 }  // namespace observers
diff --git a/bark_ml/python_wrapper/BUILD b/bark_ml/python_wrapper/BUILD
index 5cbce47d..9cbee478 100644
--- a/bark_ml/python_wrapper/BUILD
+++ b/bark_ml/python_wrapper/BUILD
@@ -1,6 +1,7 @@
 package(default_visibility = ["//visibility:public"])
 
-load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
+
+load("@pybind11_bazel//:build_defs.bzl", "pybind_library", "pybind_extension")
 
 config_setting(
   name = "op_sys",
@@ -9,11 +10,14 @@ config_setting(
 
 pybind_extension(
   name = "core",
-  srcs = glob(["*.cpp", "*.hpp"]),
+  srcs = ["main.cpp", "pyobserver.hpp", "pynn_to_value_converter.hpp"],
   deps = [
     "//bark_ml/observers:cpp_observers",
     "//bark_ml/evaluators:cpp_evaluators",
+    "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter:nn_to_value_converter",
     "//bark_ml/commons:spaces",
+    "@bark_project//bark/commons/params:params",
+    # external
     "@com_github_eigen_eigen//:eigen",
   ],
   linkopts = select({
@@ -25,3 +29,19 @@ pybind_extension(
     "//conditions:default": 1,
   })
 )
+
+pybind_library(
+  name = "polymorphic_conversion",
+  hdrs = [
+    "polymorphic_conversion.hpp",
+  ],
+  srcs=[
+    "polymorphic_conversion.cpp"
+  ],
+  deps = [
+    "@bark_project//bark/python_wrapper:common",
+    "//bark_ml/observers:cpp_observers",
+    "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter:nn_to_value_converter",
+  ],
+  visibility = ["//visibility:public"],
+)
diff --git a/bark_ml/python_wrapper/main.cpp b/bark_ml/python_wrapper/main.cpp
index b332937b..e37d2196 100644
--- a/bark_ml/python_wrapper/main.cpp
+++ b/bark_ml/python_wrapper/main.cpp
@@ -15,6 +15,7 @@
 #include "boost/variant.hpp"
 
 #include "bark/commons/params/params.hpp"
+#include "bark/commons/params/setter_params.hpp"
 #include "bark_ml/evaluators/base_evaluator.hpp"
 #include "bark_ml/evaluators/goal_reached.hpp"
 #include "bark_ml/observers/base_observer.hpp"
@@ -23,11 +24,17 @@
 #include "bark_ml/observers/static_observer.hpp"
 #include "bark_ml/evaluators/goal_reached.hpp"
 #include "bark_ml/commons/spaces.hpp"
+#include "bark_ml/python_wrapper/pyobserver.hpp"
+#include "bark_ml/python_wrapper/pynn_to_value_converter.hpp"
+#include "bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter_sequential.hpp"
+#include "bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter_policy.hpp"
+
 
 PYBIND11_DECLARE_HOLDER_TYPE(T, std::shared_ptr<T>);
 
 namespace py = pybind11;
 using bark::commons::ParamsPtr;
+using bark::commons::SetterParams;
 using bark_ml::observers::NearestObserver;
 using bark_ml::observers::FrenetObserver;
 using bark_ml::observers::StaticObserver;
@@ -35,6 +42,16 @@ using bark_ml::observers::BaseObserver;
 using bark_ml::evaluators::GoalReachedEvaluator;
 using bark_ml::spaces::Box;
 using bark_ml::spaces::Matrix_t;
+using namespace bark_ml::lib_fqf_iqn_qrdqn;
+
+py::tuple ParamsToPython(const ParamsPtr& params) {
+  return py::make_tuple(params->GetCondensedParamList());
+}
+
+ParamsPtr PythonToParams(py::tuple t) {
+  const auto param_list = t[0].cast<bark::commons::CondensedParamList>();
+  return std::make_shared<SetterParams>(true, param_list);
+}
 
 
 namespace pybind11 { namespace detail {
@@ -52,13 +69,26 @@ namespace pybind11 { namespace detail {
 }}
 
 void python_observers(py::module m) {
-  py::class_<NearestObserver,
-             std::shared_ptr<NearestObserver>>(m, "NearestObserver")
-    .def(py::init<const bark::commons::ParamsPtr&>())
+  py::class_<BaseObserver, PyObserver,
+             std::shared_ptr<BaseObserver>>(m, "BaseObserver")
+      .def(py::init<const bark::commons::ParamsPtr&>());
+
+  py::class_<NearestObserver, BaseObserver,
+              std::shared_ptr<NearestObserver>>(m, "NearestObserver")
+    .def(py::init<ParamsPtr>())
     .def("Observe", &NearestObserver::Observe)
     .def("Reset", &NearestObserver::Reset)
     .def_property_readonly(
-      "observation_space", &NearestObserver::ObservationSpace);
+      "observation_space", &NearestObserver::ObservationSpace)
+    .def(py::pickle(
+      [](const NearestObserver& o) {
+        return py::make_tuple(ParamsToPython(o.GetParams()));
+      },
+      [](py::tuple t) {
+        if (t.size() != 1)
+          throw std::runtime_error("Invalid behavior model state!");
+        return new NearestObserver(PythonToParams(t[0].cast<py::tuple>()));
+      }));
 
   py::class_<FrenetObserver,
              std::shared_ptr<FrenetObserver>>(m, "FrenetObserver")
@@ -96,6 +126,38 @@ void python_spaces(py::module m) {
     .def_property_readonly("shape", &Box<double>::shape);
 }
 
+void python_value_converters(py::module m) {
+  py::class_<NNToValueConverter, PyNNToValueConverter,
+             std::shared_ptr<NNToValueConverter>>(m, "NNToValueConverter")
+      .def(py::init<const unsigned&>());
+
+  py::class_<NNToValueConverterSequential, NNToValueConverter,
+            std::shared_ptr<NNToValueConverterSequential>>(m, "NNToValueConverterSequential")
+  .def(py::init<const unsigned&>())
+  .def(py::pickle(
+    [](const NNToValueConverterSequential& nn) {
+      return py::make_tuple(nn.GetNumActions());
+    },
+    [](py::tuple t) {
+      if (t.size() != 1)
+        throw std::runtime_error("Invalid NNToValueConverterSequential state!");
+      return new NNToValueConverterSequential(t[0].cast<unsigned>());
+    }));
+
+  py::class_<NNToValueConverterPolicy, NNToValueConverter,
+            std::shared_ptr<NNToValueConverterPolicy>>(m, "NNToValueConverterPolicy")
+  .def(py::init<const unsigned&>())
+  .def(py::pickle(
+    [](const NNToValueConverterPolicy& nn) {
+      return py::make_tuple(nn.GetNumActions());
+    },
+    [](py::tuple t) {
+      if (t.size() != 1)
+        throw std::runtime_error("Invalid NNToValueConverterPolicy state!");
+      return new NNToValueConverterPolicy(t[0].cast<unsigned>());
+    }));
+}
+
 PYBIND11_MODULE(core, m) {
   m.doc() = "Additional cpp entities for bark-ml.";
   python_observers(
@@ -104,4 +166,6 @@ PYBIND11_MODULE(core, m) {
     m.def_submodule("evaluators", "c++ evaluators"));
   python_spaces(
     m.def_submodule("spaces", "c++ spaces"));
+  python_value_converters(
+    m.def_submodule("value_converters", "value conversion"));
 }
diff --git a/bark_ml/python_wrapper/polymorphic_conversion.cpp b/bark_ml/python_wrapper/polymorphic_conversion.cpp
new file mode 100644
index 00000000..42366cd7
--- /dev/null
+++ b/bark_ml/python_wrapper/polymorphic_conversion.cpp
@@ -0,0 +1,64 @@
+// Copyright (c) 2020 Julian Bernhard
+//
+// This work is licensed under the terms of the MIT license.
+// For a copy, see <https://opensource.org/licenses/MIT>.
+
+#include "bark_ml/python_wrapper/polymorphic_conversion.hpp"
+#include "bark_ml/observers/nearest_observer.hpp"
+#include "bark/commons/params/setter_params.hpp"
+#include "bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter_sequential.hpp"
+#include "bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter_policy.hpp"
+
+namespace py = pybind11;
+
+using namespace bark_ml::observers;
+using namespace bark_ml::lib_fqf_iqn_qrdqn;
+
+py::tuple ObserverToPython(ObserverPtr observer) {
+  std::string observer_name;
+  if (typeid(*observer) == typeid(NearestObserver)) {
+    observer_name = "NearestObserver";
+  } else {
+    LOG(FATAL) << "Unknown Observer for polymorphic conversion to python: "
+               << typeid(*observer).name();
+  }
+  return py::make_tuple(observer, observer_name);
+}
+
+ObserverPtr PythonToObserver(py::tuple t) {
+  std::string observer_name = t[1].cast<std::string>();
+  if (observer_name.compare("NearestObserver") == 0) {
+    return std::make_shared<NearestObserver>(
+        t[0].cast<NearestObserver>());
+  }  else {
+       LOG(FATAL) << "Unknown Observer for polymorphic conversion to cpp: "
+               << observer_name;
+  }
+}
+
+
+py::tuple NNToValueConverterToPython(NNToValueConverterPtr converter) {
+  std::string converter_name;
+  if (typeid(*converter) == typeid(NNToValueConverterSequential)) {
+    converter_name = "NNToValueConverterSequential";
+  } else if (typeid(*converter) == typeid(NNToValueConverterPolicy)) {
+    converter_name = "NNToValueConverterPolicy";
+  } else {
+    LOG(FATAL) << "Unknown Converter for polymorphic conversion to python: "
+               << typeid(*converter).name();
+  }
+  return py::make_tuple(converter, converter_name);
+}
+NNToValueConverterPtr PythonToNNToValueConverter(py::tuple t) {
+  std::string converter_name = t[1].cast<std::string>();
+  if (converter_name.compare("NNToValueConverterSequential") == 0) {
+    return std::make_shared<NNToValueConverterSequential>(
+        t[0].cast<NNToValueConverterSequential>());
+  } else if (converter_name.compare("NNToValueConverterPolicy") == 0) {
+    return std::make_shared<NNToValueConverterPolicy>(
+        t[0].cast<NNToValueConverterPolicy>());
+  } else {
+       LOG(FATAL) << "Unknown Converter for polymorphic conversion to cpp: "
+               << converter_name;
+  }
+}
\ No newline at end of file
diff --git a/bark_ml/python_wrapper/polymorphic_conversion.hpp b/bark_ml/python_wrapper/polymorphic_conversion.hpp
new file mode 100644
index 00000000..68f58d63
--- /dev/null
+++ b/bark_ml/python_wrapper/polymorphic_conversion.hpp
@@ -0,0 +1,24 @@
+// Copyright (c) 2020 Julian Bernhard
+//
+// This work is licensed under the terms of the MIT license.
+// For a copy, see <https://opensource.org/licenses/MIT>.
+
+#ifndef BARK_OBSERVERS_POLYMORPHIC_CONVERSION_HPP_
+#define BARK_OBSERVERS_POLYMORPHIC_CONVERSION_HPP_
+
+#include "bark/python_wrapper/common.hpp"
+#include "bark_ml/observers/base_observer.hpp"
+#include "bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter.hpp"
+
+namespace py = pybind11;
+
+// For pickle we need conversion functions between the genereric base types and
+// the derived types
+
+py::tuple ObserverToPython(bark_ml::observers::ObserverPtr observer);
+bark_ml::observers::ObserverPtr PythonToObserver(py::tuple t);
+
+py::tuple NNToValueConverterToPython(bark_ml::lib_fqf_iqn_qrdqn::NNToValueConverterPtr converter);
+bark_ml::lib_fqf_iqn_qrdqn::NNToValueConverterPtr PythonToNNToValueConverter(py::tuple t);
+
+#endif  // BARK_OBSERVERS_POLYMORPHIC_CONVERSION_HPP_
diff --git a/bark_ml/python_wrapper/pynn_to_value_converter.hpp b/bark_ml/python_wrapper/pynn_to_value_converter.hpp
new file mode 100644
index 00000000..a5c98904
--- /dev/null
+++ b/bark_ml/python_wrapper/pynn_to_value_converter.hpp
@@ -0,0 +1,32 @@
+// Copyright (c) 2020 fortiss GmbH
+//
+// Authors: Julian Bernhard, Klemens Esterle, Patrick Hart and
+// Tobias Kessler
+//
+// This work is licensed under the terms of the MIT license.
+// For a copy, see <https://opensource.org/licenses/MIT>.
+
+#ifndef BARK_ML_PYTHON_WRAPPER_PY_NN_TO_VALUE_CONVERTER_HPP_
+#define BARK_ML_PYTHON_WRAPPER_PY_NN_TO_VALUE_CONVERTER_HPP_
+#include <memory>
+
+#include "bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model/nn_to_value_converter/nn_to_value_converter.hpp"
+
+namespace py = pybind11;
+using bark_ml::lib_fqf_iqn_qrdqn::NNToValueConverter;
+using bark_ml::lib_fqf_iqn_qrdqn::ValueType;
+using ValueMap = std::unordered_map<ValueType, std::vector<double>>;
+
+class PyNNToValueConverter : public NNToValueConverter {
+ public:
+  using NNToValueConverter::NNToValueConverter;
+
+  ValueMap ConvertToValueMap(
+                const std::vector<float>& nn_output) const {
+  PYBIND11_OVERLOAD_PURE(ValueMap, NNToValueConverter, ConvertToValueMap,
+                           nn_output); }
+
+};
+
+
+#endif  // BARK_ML_PYTHON_WRAPPER_PY_NN_TO_VALUE_CONVERTER_HPP_
diff --git a/bark_ml/python_wrapper/pyobserver.hpp b/bark_ml/python_wrapper/pyobserver.hpp
new file mode 100644
index 00000000..afff8698
--- /dev/null
+++ b/bark_ml/python_wrapper/pyobserver.hpp
@@ -0,0 +1,42 @@
+// Copyright (c) 2020 fortiss GmbH
+//
+// Authors: Julian Bernhard, Klemens Esterle, Patrick Hart and
+// Tobias Kessler
+//
+// This work is licensed under the terms of the MIT license.
+// For a copy, see <https://opensource.org/licenses/MIT>.
+
+#ifndef BARK_ML_PYTHON_WRAPPER_PYOBSERVER_HPP_
+#define BARK_ML_PYTHON_WRAPPER_PYOBSERVER_HPP_
+#include <memory>
+
+#include "bark_ml/observers/base_observer.hpp"
+
+namespace py = pybind11;
+using bark_ml::observers::BaseObserver;
+using bark_ml::spaces::Box;
+using bark::world::WorldPtr;
+using bark::world::ObservedWorldPtr;
+using bark::world::ObservedWorld;
+using ObservedState = Eigen::Matrix<double, 1, Eigen::Dynamic>;
+
+class PyObserver : public BaseObserver {
+ public:
+  using BaseObserver::BaseObserver;
+
+  ObservedState Observe(
+    const ObservedWorld& observed_world) const {
+    PYBIND11_OVERLOAD_PURE(ObservedState, BaseObserver, Observe,
+                           observed_world); }
+
+  WorldPtr Reset(const WorldPtr& world) {
+    PYBIND11_OVERLOAD_PURE(WorldPtr, BaseObserver, Reset,
+                           world); }
+  Box<double> ObservationSpace() const {
+    PYBIND11_OVERLOAD_PURE(Box<double>, BaseObserver, ObservationSpace,
+                           ); }
+
+};
+
+
+#endif  // BARK_ML_PYTHON_WRAPPER_PYOBSERVER_HPP_
diff --git a/examples/BUILD b/examples/BUILD
index fd5b3c89..fcace66f 100644
--- a/examples/BUILD
+++ b/examples/BUILD
@@ -123,19 +123,3 @@ py_test(
   ]
 )
 
-py_test(
-  name = "diadem_dqn",
-  srcs = ["diadem_dqn.py"],
-  data = ["@bark_project//bark/python_wrapper:core.so",
-          "//bark_ml:generate_core",
-          "//examples/example_params:example_params"
-        ],
-  imports = ["../bark_ml/python_wrapper/",
-             "./external/diadem_project/"],
-  deps = [
-    "//bark_ml/environments:single_agent_runtime",
-    "//bark_ml/environments:blueprints",
-    "//bark_ml/observers:observers",
-    "//bark_ml/library_wrappers/lib_diadem:diadem"
-  ]
-)
diff --git a/examples/example_params/fqf_params.json b/examples/example_params/fqf_params.json
index f6e7e644..ac34dfc7 100644
--- a/examples/example_params/fqf_params.json
+++ b/examples/example_params/fqf_params.json
@@ -78,29 +78,6 @@
             "Seed":0
         }
     },
-    "BehaviorIDMClassic":{
-        "MinimumSpacing":2.0,
-        "DesiredTimeHeadway":1.5,
-        "MaxAcceleration":1.7000000476837158,
-        "AccelerationLowerBound":-5.0,
-        "AccelerationUpperBound":8.0,
-        "DesiredVelocity":15.0,
-        "ComfortableBrakingAcceleration":1.6699999570846558,
-        "MinVelocity":0.0,
-        "MaxVelocity":50.0,
-        "Exponent":4,
-        "BrakeForLaneEnd":false,
-        "BrakeForLaneEndEnabledDistance":60.0,
-        "BrakeForLaneEndDistanceOffset":15.0,
-        "NumTrajectoryTimePoints":11.0
-    },
-    "DynamicModel":{
-        "wheel_base":2.700000047683716,
-        "delta_max":0.20000000298023224,
-        "lat_acc_max":4.0,
-        "lon_acceleration_max":4.0,
-        "lon_acceleration_min":-8.0
-    },
     "agent":{
         "MaxHistoryLength":50
     },
diff --git a/utils/dependencies.bzl b/utils/dependencies.bzl
index c8503729..3bf9fefb 100644
--- a/utils/dependencies.bzl
+++ b/utils/dependencies.bzl
@@ -13,54 +13,29 @@ def bark_ml_dependencies():
     branch = "master",
     remote = "https://github.com/bark-simulator/bark",
   )
-  # _maybe(
-  #   native.local_repository,
-  #   name = "bark_project",
-  #   path = "/home/xliu/bark"
-  # )
 
   _maybe(
-      git_repository,
-      name = "pybind11_bazel",
-      commit="c4a29062b77bf42836d995f6ce802f642cffb939",
-      remote = "https://github.com/bark-simulator/pybind11_bazel"
-  )
-
-  # _maybe(
-  #   native.local_repository,
-  #   name = "pybind11_bazel",
-  #   path = "/Users/hart/Development/pybind11_bazel"
-  # )
-  # alternative to torch api used from virtual env
-  _maybe(
-    http_archive,
-    name = "torch_api",
-    urls = ["https://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip"],
+    native.new_local_repository,
+    name = "torchcpp",
+    path = "./bark_ml/python_wrapper/venv/lib/python3.7/site-packages/",
     build_file_content = """
 cc_library(
     name = "lib",
-    srcs = glob(["libtorch/lib/*.*"]),
-    hdrs = glob(["libtorch/include/**/*.h", "libtorch/include/*.h"]),
+    srcs = ["torch/lib/libtorch.so",
+                 "torch/lib/libc10.so", "torch/lib/libtorch_cpu.so"],
+    hdrs = glob(["torch/include/**/*.h", "torch/include/*.h"]),
     visibility = ["//visibility:public"],
-)
-    """)
-
-#   _maybe(
-#     native.new_local_repository,
-#     name = "torchcpp",
-#     path = "./bark_ml/python_wrapper/venv/lib/python3.7/site-packages/",
-#     build_file_content = """
-# cc_library(
-#     name = "lib",
-#     srcs = ["torch/lib/libc10.so", "torch/lib/libtorch_cpu.so"],
-#     hdrs = glob(["torch/include/**/*.h", "torch/include/*.h"]),
-#     visibility = ["//visibility:public"],
-# )
-#     """)
+    linkopts = [
+        "-ltorch",
+        "-ltorch_cpu",
+        "-lc10",
+    ],
+)"""
+  )
 
-#   _maybe(
-#     git_repository,
-#     name = "diadem_project",
-#     commit = "64b2987fbdd69ad533f30b545568c691ad5afb00",
-#     remote = "https://github.com/juloberno/diadem"
-#   )
\ No newline at end of file
+  _maybe(
+      git_repository,
+      name = "pybind11_bazel",
+      commit="c4a29062b77bf42836d995f6ce802f642cffb939",
+      remote = "https://github.com/bark-simulator/pybind11_bazel"
+  )
diff --git a/utils/dev_into.sh b/utils/dev_into.sh
index baabd96d..7ffeb43c 100644
--- a/utils/dev_into.sh
+++ b/utils/dev_into.sh
@@ -1,2 +1,5 @@
 #!/bin/bash
-source ./bark_ml/python_wrapper/venv/bin/activate 
\ No newline at end of file
+DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+PARENTDIR=$(dirname "$DIR")
+source $PARENTDIR/bark_ml/python_wrapper/venv/bin/activate
+export LD_LIBRARY_PATH=$PARENTDIR/bark_ml/python_wrapper/venv/lib/python3.7/site-packages/torch/lib/ 
diff --git a/utils/diadem_install.sh b/utils/diadem_install.sh
deleted file mode 100644
index a0c23890..00000000
--- a/utils/diadem_install.sh
+++ /dev/null
@@ -1,3 +0,0 @@
-#!/bin/bash
-virtualenv -p python3.7 ./python/venv
-source ./bark_ml/python_wrapper/venv/bin/activate && pip3.7 install --no-cache-dir --upgrade --trusted-host pypi.org -r utils/docker/installers/requirements_diadem.txt 
\ No newline at end of file
diff --git a/utils/docker/installers/requirements.txt b/utils/docker/installers/requirements.txt
index 96cd6b48..44d816c0 100644
--- a/utils/docker/installers/requirements.txt
+++ b/utils/docker/installers/requirements.txt
@@ -1,4 +1,3 @@
-
 aabbtree==2.8.1
 absl-py==0.13.0
 lxml==4.6.3
@@ -14,6 +13,9 @@ tensorflow==2.5.0
 tensorflow-addons==0.13.0
 tensorflow-estimator==2.5.0
 tensorflow-probability==0.11.0
+tf2_gnn>=2.4.0
+spektral>=0.6.0
+ray==0.8.5
 tensorboard==2.5.0
 tf-agents==0.6.0
 notebook==6.4.0
diff --git a/utils/docker/installers/requirements_diadem.txt b/utils/docker/installers/requirements_diadem.txt
deleted file mode 100644
index d14cb149..00000000
--- a/utils/docker/installers/requirements_diadem.txt
+++ /dev/null
@@ -1,22 +0,0 @@
-tensorflow==1.15.5
-numpy==1.18.1
-pyyaml
-matplotlib==3.0.3
-lxml>=4.4.2
-scipy==1.4.1
-sphinx==2.3.1
-recommonmark==0.6.0
-sphinx_rtd_theme==0.4.3
-pandas==0.24.2
-autopep8==1.4.4
-cpplint==1.4.4
-pygame==1.9.6
-panda3d==1.10.5
-aabbtree==2.3.1
-ray==0.7.6
-psutil==5.6.6
-notebook>=6.0.3
-jupyter==1.0.0
-gym==0.15.4
-ipython==7.13.0
-svgutils==0.3.1
\ No newline at end of file
diff --git a/utils/docker/installers/requirements_diadem_cluster.txt b/utils/docker/installers/requirements_diadem_cluster.txt
deleted file mode 100644
index 37034bb0..00000000
--- a/utils/docker/installers/requirements_diadem_cluster.txt
+++ /dev/null
@@ -1,22 +0,0 @@
-tensorflow-gpu==1.15.5
-numpy==1.18.1
-pyyaml
-matplotlib==3.0.3
-lxml>=4.4.2
-scipy==1.4.1
-sphinx==2.3.1
-recommonmark==0.6.0
-sphinx_rtd_theme==0.4.3
-pandas==0.24.2
-autopep8==1.4.4
-cpplint==1.4.4
-pygame==1.9.6
-panda3d==1.10.5
-aabbtree==2.3.1
-ray==0.7.6
-psutil==5.6.6
-notebook>=6.0.3
-jupyter==1.0.0
-gym==0.15.4
-ipython==7.13.0
-svgutils==0.3.1
\ No newline at end of file
diff --git a/utils/install.sh b/utils/install.sh
index 915e4b65..0046ebbe 100644
--- a/utils/install.sh
+++ b/utils/install.sh
@@ -1,3 +1,5 @@
 #!/bin/bash
 virtualenv -p python3.7 ./bark_ml/python_wrapper/venv
-source ./bark_ml/python_wrapper/venv/bin/activate && pip install --no-cache-dir --upgrade --trusted-host pypi.org -r ./utils/docker/installers/requirements.txt 
\ No newline at end of file
+source ./bark_ml/python_wrapper/venv/bin/activate && pip install --no-cache-dir --upgrade --trusted-host pypi.org -r ./utils/docker/installers/requirements.txt 
+# cpu version
+pip install torch==1.9.0+cpu torchvision -f https://download.pytorch.org/whl/torch_stable.html
\ No newline at end of file


Test patch:
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/BUILD b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/BUILD
new file mode 100644
index 00000000..68b43a3a
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/BUILD
@@ -0,0 +1,93 @@
+py_test(
+    name = "save_load_test",
+    srcs = ["save_load_test.py"],
+    data = [
+            "@bark_project//bark:generate_core"
+            ],
+    imports = ["../external/bark_project/bark/python_wrapper/",
+              ],
+    deps = ["//bark_ml/environments:single_agent_runtime",
+            "//bark_ml/environments:gym",
+            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent:agents",
+            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations:demonstrations",
+            "//bark_ml/commons:py_spaces"
+           ],
+    visibility = ["//visibility:public"],
+)
+
+py_library(
+   name = "test_demo_behavior",
+   srcs = ["test_demo_behavior.py"]
+)
+
+py_test(
+    name = "demonstration_collector_test",
+    srcs = [
+            "demonstration_collector_test.py"
+            ],
+    data = [
+            "@bark_project//bark:generate_core"
+            ],
+    imports = ["../external/bark_project/bark/python_wrapper/",
+              ],
+    deps = ["//bark_ml/environments:single_agent_runtime",
+            "//bark_ml/environments:gym",
+            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent:agents",
+            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations:demonstrations",
+            "//bark_ml/commons:py_spaces",
+            ":test_demo_behavior"
+           ],
+    visibility = ["//visibility:public"],
+)
+
+py_test(
+    name = "model_loader_tests",
+    srcs = ["model_loader_tests.py",
+            "test_imitation_agent.py"
+            ],
+    data = ["@torchcpp//:lib",
+            "@bark_project//bark/python_wrapper:core.so",
+            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/lib_fqf_iqn_qrdqn_test_data:weights",
+            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_loader:pytorch_script_wrapper.so"],
+    imports = ["../external/bark_project/bark/python_wrapper/",
+               "../external/torchcpp/torch/lib" # this path does not exist even after adding data dependency. Keeping it for further research. 
+                                                # and ../ should be ../../ for bark as well as pytorch as there is no external 
+                                                # in runfiles/bark_ml/bark_ml/external (not exists but the import adds it!)
+              ],
+    deps = ["//bark_ml/environments:single_agent_runtime",
+            "//bark_ml/environments:gym",
+            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent:agents",
+            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations:demonstrations",
+            "//bark_ml/commons:py_spaces"
+           ],
+    visibility = ["//visibility:public"],
+)
+
+py_test(
+    name = "test_imitation_agent",
+    srcs = ["test_imitation_agent.py"],
+    data = [
+        "//bark_ml:generate_core",
+        "@bark_project//bark:generate_core",
+    ],
+    deps = [
+        "//bark_ml/environments:blueprints",
+        "//bark_ml/environments:gym",
+        "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent:agents",
+        "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent/demonstrations:demonstrations",
+        "@bark_project//bark/runtime/viewer:py_viewer",
+        "@benchmark_database//load:benchmark_database",
+        "@benchmark_database//serialization:database_serializer"
+    ],
+    visibility = ["//visibility:public"],
+)
+
+test_suite(
+  name = "py_lib_fqf_imitation_agent_tests",
+  tests = [
+    ":save_load_test",
+    ":demonstration_collector_test",
+    ":model_loader_tests",
+    ":test_imitation_agent"
+  ]
+)
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/demonstration_collector_test.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/demonstration_collector_test.py
new file mode 100644
index 00000000..1da736e0
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/demonstration_collector_test.py
@@ -0,0 +1,107 @@
+# Copyright (c) 2020 fortiss GmbH
+#
+# Authors: Julian Bernhard, Patrick Hart
+#
+# This work is licensed under the terms of the MIT license.
+# For a copy, see <https://opensource.org/licenses/MIT>.
+
+try:
+    import debug_settings
+except:
+    pass
+
+
+import unittest
+import numpy as np
+import os
+import gym
+import matplotlib
+import time
+
+# BARK imports
+from bark.runtime.commons.parameters import ParameterServer
+
+# BARK-ML imports
+from bark_ml.environments.blueprints import \
+  DiscreteHighwayBlueprint, DiscreteMergingBlueprint
+from bark_ml.environments.single_agent_runtime import SingleAgentRuntime
+import bark_ml.environments.gym
+from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.agent.demonstrations import DemonstrationCollector
+from bark_ml.observers.nearest_state_observer import NearestAgentsObserver
+from bark_ml.behaviors.discrete_behavior import BehaviorDiscreteMacroActionsML
+
+import bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.tests.test_demo_behavior
+
+class TestEvaluator:
+  reach_goal = True
+  def __init__(self,
+               params=ParameterServer()):
+    self.step = 0
+
+  def _evaluate(self, observed_world, eval_results, action):
+    """Returns information about the current world state
+    """
+    done = False
+    reward = 0.0
+    info = {"goal_r1" : False}
+    if self.step > 2:        
+      done = True
+      if self.reach_goal:
+        reward = 0.1
+        info = {"goal_r1" : True}
+    self.step += 1
+    return reward, done, info
+
+  def Evaluate(self, observed_world, action):
+    """Evaluates the observed world
+    """
+    eval_results, reward, done = None, 0., False
+    eval_results = observed_world.Evaluate()
+    reward, done, eval_results = self._evaluate(
+      observed_world, eval_results, action)
+    return reward, done, eval_results
+
+  def Reset(self, world):
+    world.ClearEvaluators()
+    self._add_evaluators()
+    for key, evaluator in self._evaluators.items():
+      world.AddEvaluator(key, evaluator)
+    return world
+    
+  def Reset(self, world):
+    self._step = 0
+    #every second scenario goal is not reached
+    TestEvaluator.reach_goal = not TestEvaluator.reach_goal 
+
+
+class DemonstrationCollectorTests(unittest.TestCase):
+  def test_collect_demonstrations(self):
+    params = ParameterServer()
+    bp = DiscreteHighwayBlueprint(params, num_scenarios=10, random_seed=0)
+    env = SingleAgentRuntime(blueprint=bp, render=False)
+    env._observer = NearestAgentsObserver(params)
+    env._action_wrapper = BehaviorDiscreteMacroActionsML(params)
+    env._evaluator = TestEvaluator()
+    
+    demo_behavior = bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.\
+            tests.test_demo_behavior.TestDemoBehavior(params)
+    collector = DemonstrationCollector()
+    collection_result = collector.CollectDemonstrations(4, "./test_demo_collected", env=env, motion_primitive_behavior=demo_behavior, \
+           use_mp_runner=False, runner_init_params={"deepcopy" : False})
+    self.assertTrue(os.path.exists("./test_demo_collected/collection_result"))
+    print(collection_result.get_data_frame().to_string())
+
+    experiences = collector.ProcessCollectionResult(eval_criteria = {"goal_r1" : lambda x : x})
+    # expected length = 2 scenarios (only every second reaches goal) x 3 steps (4 executed, but first not counted)
+    self.assertEqual(len(experiences), 2*3) 
+
+    collector.dump("./final_collections")
+
+    loaded_collector = DemonstrationCollector.load("./final_collections")
+    experiences_loaded = loaded_collector.GetDemonstrationExperiences()
+    print(experiences_loaded)
+    self.assertEqual(len(experiences_loaded), 2*3) 
+
+
+if __name__ == '__main__':
+  unittest.main()
diff --git a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/BUILD b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/lib_fqf_iqn_qrdqn_test_data/BUILD
similarity index 100%
rename from bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/BUILD
rename to bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/lib_fqf_iqn_qrdqn_test_data/BUILD
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/model_loader_tests.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/model_loader_tests.py
new file mode 100644
index 00000000..3a6ba582
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/model_loader_tests.py
@@ -0,0 +1,112 @@
+# Copyright (c) 2020 fortiss GmbH
+#
+# Authors: Patrick Hart
+#
+# This work is licensed under the terms of the MIT license.
+# For a copy, see <https://opensource.org/licenses/MIT>.
+
+try:
+    import debug_settings
+except:
+    pass
+
+import unittest
+import numpy as np
+import os
+import gym
+import time
+import logging
+
+# BARK imports
+from bark.runtime.commons.parameters import ParameterServer
+
+from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.model_loader \
+ import pytorch_script_wrapper
+
+from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.tests.test_imitation_agent import TestActionWrapper, \
+        TestObserver, create_data
+
+from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.agent import ImitationAgent
+num_actions = 4
+class TestMotionPrimitiveBehavior:
+  def __init__(self, num_actions):
+    self._num_actions = num_actions
+
+  def GetMotionPrimitives(self):
+    return list(range(0,self._num_actions))
+
+class TestDemonstrationCollector:
+  def __init__(self):
+    self.data = create_data(10000)
+    self._observer = TestObserver()
+    self._ml_behavior =  TestActionWrapper()
+    self.motion_primitive_behavior = TestMotionPrimitiveBehavior(num_actions)
+
+  def GetDemonstrationExperiences(self):
+    return self.data
+
+  @property
+  def observer(self):
+    return self._observer
+
+  @property
+  def ml_behavior(self):
+    return self._ml_behavior
+
+  def GetDirectory(self):
+    return "./save_dir/collections"
+
+def imitation_agent(layer_dims):
+    params = ParameterServer()
+    params["ML"]["BaseAgent"]["NumSteps"] = 2
+    params["ML"]["BaseAgent"]["EvalInterval"] = 1
+    params["ML"]["ImitationModel"]["EmbeddingDims"] = layer_dims
+    data = create_data(1000)
+    demo_train = data[0:1000]
+    demo_test = data[51:]
+    agent = ImitationAgent(agent_save_dir="./save_dir", demonstration_collector=TestDemonstrationCollector(),
+                          params=params)
+    agent.run()
+    return agent, demo_train, demo_test
+
+
+class ModelLoaderTests(unittest.TestCase):
+  def test_model_loader_imitation(self):
+    agent, demo_train, demo_test = imitation_agent([200, 200, 100])
+    agent.save(checkpoint_type="last")
+    script_file = agent.get_script_filename("last")
+    # test all network
+    model = pytorch_script_wrapper.ModelLoader()
+    model.LoadModel(os.path.abspath(script_file))
+
+    num_iters = 10  # Number of times to repeat experiment to calcualte runtime
+
+    # Time num_iters iterations for inference using C++ model
+    start = time.time()
+    output_cpp = []
+    for idx in range(num_iters):
+      output_cpp.append(model.Inference(list(demo_train[idx][0])))
+    end = time.time()
+    time_cpp = end - start  # todo - how to analyze python vs cpp test time in tests?
+    logging.info(f"C++ time: {time_cpp/num_iters}")
+
+    #Time num_iters iterations for inference using python model
+    start = time.time()
+    output_py = []
+    agent.online_net.eval()
+    for idx in range(num_iters):
+      output_py.append(agent.calculate_actions(demo_train[idx][0]))
+
+    end = time.time()
+    time_py = end - start
+    logging.info(f"Python time: {time_py/num_iters}")
+
+    np.testing.assert_array_almost_equal(
+        np.asarray(output_cpp),
+        np.asarray(output_py),
+        decimal=6,
+        err_msg="C++ and python models don't match")
+
+
+if __name__ == '__main__':
+  unittest.main()
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/save_load_test.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/save_load_test.py
new file mode 100644
index 00000000..5a5e85e8
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/save_load_test.py
@@ -0,0 +1,67 @@
+# Copyright (c) 2020 fortiss GmbH
+#
+# Authors: Julian Bernhard, Patrick Hart
+#
+# This work is licensed under the terms of the MIT license.
+# For a copy, see <https://opensource.org/licenses/MIT>.
+
+try:
+    import debug_settings
+except:
+    pass
+
+
+import unittest
+from gym import spaces
+import numpy as np
+
+# BARK imports
+from bark.runtime.commons.parameters import ParameterServer
+
+# BARK-ML imports
+from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.agent import FQFAgent
+
+observation_length = 5 
+num_actions = 4
+
+class TestObserver:
+  @property
+  def observation_space(self):
+    # TODO(@hart): use from spaces.py
+    return spaces.Box(
+      low=np.zeros(observation_length),
+      high = np.ones(observation_length))
+
+class TestActionWrapper():
+  @property
+  def action_space(self):
+    return spaces.Discrete(num_actions)
+
+class DummyEnv:
+  def __init__(self):
+    self._observer = TestObserver()
+    self._ml_behavior =  TestActionWrapper()
+
+  def reset(self):
+    pass
+
+  def step(self):
+    pass
+
+class BaseAgentTests(unittest.TestCase):
+  def test_agents(self):
+    params = ParameterServer()
+
+    fqf_agent = FQFAgent(env = DummyEnv(), agent_save_dir="./save_dir", params=params)
+
+    fqf_agent.save(checkpoint_type="best")
+    fqf_agent.save(checkpoint_type="last")
+
+    loaded_agent = FQFAgent(env = DummyEnv(), agent_save_dir="./save_dir", checkpoint_load="best")
+
+    self.assertEqual(loaded_agent.ml_behavior.action_space.n, fqf_agent.ml_behavior.action_space.n)
+    self.assertEqual(loaded_agent.ent_coef, fqf_agent.ent_coef)
+
+
+if __name__ == '__main__':
+  unittest.main()
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/test_demo_behavior.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/test_demo_behavior.py
new file mode 100644
index 00000000..dd0159d0
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/test_demo_behavior.py
@@ -0,0 +1,39 @@
+# Copyright (c) 2020 fortiss GmbH
+#
+# Authors: Julian Bernhard, Patrick Hart
+#
+# This work is licensed under the terms of the MIT license.
+# For a copy, see <https://opensource.org/licenses/MIT>.
+
+import numpy as np
+from bark.core.models.behavior import BehaviorModel
+
+class TestDemoBehavior(BehaviorModel):
+  """Dummy Python behavior model
+  """
+  def __init__(self,
+               dynamic_model = None,
+               params = None):
+    BehaviorModel.__init__(self, params)
+    self._dynamic_model = dynamic_model
+    self._params = params
+
+  def Plan(self, delta_time, world):
+    super(TestDemoBehavior, self).ActionToBehavior(
+      np.array([2., 1.], dtype=np.float32))
+    trajectory = np.array([[0., 0., 0., 0., 0.],
+                           [0., 0., 0., 0., 0.]], dtype=np.float32)
+    super(TestDemoBehavior, self).SetLastTrajectory(trajectory)
+    return trajectory
+
+  def Clone(self):
+    return self
+
+  def GetLastMacroAction(self):
+    return 22
+
+  def __setstate__(self, d):
+    pass
+
+  def __getstate__(self):
+    return {}
\ No newline at end of file
diff --git a/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/test_imitation_agent.py b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/test_imitation_agent.py
new file mode 100644
index 00000000..176d1cae
--- /dev/null
+++ b/bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/tests/test_imitation_agent.py
@@ -0,0 +1,102 @@
+# Copyright (c) 2020 fortiss GmbH
+#
+# Authors: Julian Bernhard, Patrick Hart
+#
+# This work is licensed under the terms of the MIT license.
+# For a copy, see <https://opensource.org/licenses/MIT>.
+
+try:
+    import debug_settings
+except:
+    pass
+
+import unittest
+import numpy as np
+from gym import spaces
+
+# BARK imports
+from bark.runtime.commons.parameters import ParameterServer
+from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.agent import ImitationAgent
+
+observation_length = 5 
+num_actions = 4
+
+class TestObserver:
+  @property
+  def observation_space(self):
+    # TODO(@hart): use from spaces.py
+    return spaces.Box(
+      low=np.zeros(observation_length),
+      high = np.ones(observation_length))
+
+class TestActionWrapper():
+  @property
+  def action_space(self):
+    return spaces.Discrete(num_actions)
+
+def action_values_at_state(state):
+  envelope_costs = []
+  collision_costs = []
+  return_values = []
+  for idx in range(1, num_actions+1):
+    envelope_costs.append(state[0]* 1.0/idx + state[1]* 1.0/idx)
+    collision_costs.append(state[1]* 1.0/idx*0.001 + state[3]* 1.0/idx*0.001)
+    collision_costs.append(state[3]* 1.0/idx*0.2 + state[4]* 1.0/idx*0.1)
+  action_values = []
+  action_values.extend(envelope_costs)
+  action_values.extend(collision_costs)
+  action_values.extend(return_values)
+  return [state, action_values]
+
+def create_data(num):
+  observations = np.random.rand(num, observation_length)
+  action_values_data = np.apply_along_axis(action_values_at_state, 1, observations)
+  return action_values_data
+
+class TestMotionPrimitiveBehavior:
+  def __init__(self, num_actions):
+    self._num_actions = num_actions
+
+  def GetMotionPrimitives(self):
+    return list(range(0,self._num_actions))
+
+class TestDemonstrationCollector:
+  def __init__(self):
+    self.data = create_data(10000)
+    self._observer = TestObserver()
+    self._ml_behavior =  TestActionWrapper()
+    self.motion_primitive_behavior = TestMotionPrimitiveBehavior(num_actions)
+
+  def GetDemonstrationExperiences(self):
+    return self.data
+
+  @property
+  def observer(self):
+    return self._observer
+
+  @property
+  def ml_behavior(self):
+    return self._ml_behavior
+
+  def GetDirectory(self):
+    return ""
+
+
+class EvaluationTests(unittest.TestCase):
+  # make sure the agent works
+  def test_agent_wrapping(self):
+    params = ParameterServer()
+    params["ML"]["BaseAgent"]["NumSteps"] = 10
+    params["ML"]["BaseAgent"]["MaxEpisodeSteps"] = 2
+    params["ML"]["BaseAgent"]["EvalInterval"] = 2
+    params["ML"]["BaseAgent"]["TrainTestRatio"] = 0.2
+    
+    demonstration_collector = TestDemonstrationCollector()
+
+    agent = ImitationAgent(agent_save_dir="./save_dir", 
+                           demonstration_collector=demonstration_collector,
+                           params=params, checkpoint_load = False)
+    agent.run()
+
+if __name__ == '__main__':
+  unittest.main()
diff --git a/bark_ml/tests/BUILD b/bark_ml/tests/BUILD
index 8a98d2c5..1f913d07 100644
--- a/bark_ml/tests/BUILD
+++ b/bark_ml/tests/BUILD
@@ -84,26 +84,6 @@ py_test(
     visibility = ["//visibility:public"],
 )
 
-py_test(
-    name = "py_library_torch_agents_tests",
-    srcs = ["py_library_torch_agents_tests.py"],
-    data = ["@torchcpp//:lib",
-            "@bark_project//bark/python_wrapper:core.so",
-            "//bark_ml/tests/lib_fqf_iqn_qrdqn_test_data:weights",
-            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/model_wrapper:pytorch_script_wrapper.so"],
-    imports = ["../external/bark_project/bark/python_wrapper/",
-               "../external/torchcpp/torch/lib" # this path does not exist even after adding data dependency. Keeping it for further research.
-                                                # and ../ should be ../../ for bark as well as pytorch as there is no external
-                                                # in runfiles/bark_ml/bark_ml/external (not exists but the import adds it!)
-              ],
-    deps = ["//bark_ml/environments:single_agent_runtime",
-            "//bark_ml/environments:gym",
-            "//bark_ml/library_wrappers/lib_fqf_iqn_qrdqn/agent:agents",
-            "//bark_ml/commons:py_spaces"
-           ],
-    visibility = ["//visibility:public"],
-)
-
 py_test(
     name = "py_bark_behavior_model_tests",
     srcs = ["py_bark_behavior_model_tests.py"],
@@ -128,26 +108,33 @@ py_test(
     visibility = ["//visibility:public"],
 )
 
-py_test(
-    name = "py_graph_nets_tests",
-    srcs = ["py_graph_nets_tests.py"],
-    data = ["@bark_project//bark/python_wrapper:core.so",
-            "//bark_ml:generate_core"],
-    imports = ["../external/bark_project/bark/python_wrapper/"],
+
+cc_test(
+    name = "observer_nearest_test",
+    srcs = [
+        "observer_nearest_test.cc",
+    ],
+    copts = ["-Iexternal/gtest/include"],
     deps = [
-        "//bark_ml/environments:single_agent_runtime",
-        "//bark_ml/environments:blueprints",
-        "//bark_ml/library_wrappers/lib_tf_agents/networks:networks",
-        "//bark_ml/library_wrappers/lib_tf_agents/agents:agents",
-        "//bark_ml/library_wrappers/lib_tf_agents/runners:runners",
+        "@bark_project//bark/geometry",
+        "@bark_project//bark/world:world",
+        "@bark_project//bark/world/map:map_interface",
+        "@bark_project//bark/world/map:roadgraph",
+        "@bark_project//bark/world/opendrive:opendrive",
+        "@bark_project//bark/models/behavior/constant_acceleration:constant_acceleration",
+        "@bark_project//bark/models/execution/interpolation:interpolation",
+        "@bark_project//bark/world/evaluation:evaluation",
+        "@bark_project//bark/models/behavior/motion_primitives:motion_primitives",
+        "@bark_project//bark/world/tests:make_test_world",
+        "//bark_ml/observers:cpp_observers",
+        "@gtest//:gtest_main",
     ],
-    visibility = ["//visibility:public"],
 )
 
 cc_test(
-    name = "observer_test",
+    name = "observer_static_test",
     srcs = [
-        "observer_test.cc",
+        "observer_static_test.cc",
     ],
     copts = ["-Iexternal/gtest/include"],
     deps = [
@@ -164,4 +151,4 @@ cc_test(
         "//bark_ml/observers:cpp_observers",
         "@gtest//:gtest_main",
     ],
-)
\ No newline at end of file
+)
diff --git a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/fqf/online_net.pth b/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/fqf/online_net.pth
deleted file mode 100644
index 8ea5a730..00000000
Binary files a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/fqf/online_net.pth and /dev/null differ
diff --git a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/fqf/online_net_script.pt b/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/fqf/online_net_script.pt
deleted file mode 100644
index cbb33321..00000000
Binary files a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/fqf/online_net_script.pt and /dev/null differ
diff --git a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/fqf/target_net.pth b/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/fqf/target_net.pth
deleted file mode 100644
index eff35243..00000000
Binary files a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/fqf/target_net.pth and /dev/null differ
diff --git a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/iqn/online_net.pth b/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/iqn/online_net.pth
deleted file mode 100644
index 45ad1d40..00000000
Binary files a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/iqn/online_net.pth and /dev/null differ
diff --git a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/iqn/online_net_script.pt b/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/iqn/online_net_script.pt
deleted file mode 100644
index 9db07070..00000000
Binary files a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/iqn/online_net_script.pt and /dev/null differ
diff --git a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/iqn/target_net.pth b/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/iqn/target_net.pth
deleted file mode 100644
index 45728193..00000000
Binary files a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/iqn/target_net.pth and /dev/null differ
diff --git a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/qrdqn/online_net.pth b/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/qrdqn/online_net.pth
deleted file mode 100644
index 816a292c..00000000
Binary files a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/qrdqn/online_net.pth and /dev/null differ
diff --git a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/qrdqn/online_net_script.pt b/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/qrdqn/online_net_script.pt
deleted file mode 100644
index 04994e68..00000000
Binary files a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/qrdqn/online_net_script.pt and /dev/null differ
diff --git a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/qrdqn/target_net.pth b/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/qrdqn/target_net.pth
deleted file mode 100644
index 954fe7d8..00000000
Binary files a/bark_ml/tests/lib_fqf_iqn_qrdqn_test_data/qrdqn/target_net.pth and /dev/null differ
diff --git a/bark_ml/tests/observer_nearest_test.cc b/bark_ml/tests/observer_nearest_test.cc
new file mode 100644
index 00000000..71492843
--- /dev/null
+++ b/bark_ml/tests/observer_nearest_test.cc
@@ -0,0 +1,208 @@
+// Copyright (c) 2020 fortiss GmbH
+//
+// Authors: Julian Bernhard, Klemens Esterle, Patrick Hart and
+// Tobias Kessler
+//
+// This work is licensed under the terms of the MIT license.
+// For a copy, see <https://opensource.org/licenses/MIT>.
+
+#include "bark/world/observed_world.hpp"
+#include "bark/commons/params/setter_params.hpp"
+#include "bark/geometry/polygon.hpp"
+#include "bark/geometry/standard_shapes.hpp"
+#include "bark/models/behavior/constant_acceleration/constant_acceleration.hpp"
+#include "bark/models/behavior/motion_primitives/continuous_actions.hpp"
+#include "bark/models/dynamic/single_track.hpp"
+#include "bark/models/dynamic/single_track_steering_rate.hpp"
+#include "bark/models/execution/interpolation/interpolate.hpp"
+#include "bark/world/evaluation/evaluator_collision_agents.hpp"
+#include "bark/world/goal_definition/goal_definition.hpp"
+#include "bark/world/goal_definition/goal_definition_polygon.hpp"
+#include "bark/world/map/map_interface.hpp"
+#include "bark/world/map/roadgraph.hpp"
+#include "bark/world/objects/agent.hpp"
+#include "bark/world/opendrive/opendrive.hpp"
+#include "bark/world/tests/make_test_world.hpp"
+#include "bark/world/tests/make_test_xodr_map.hpp"
+#include "bark_ml/observers/nearest_observer.hpp"
+#include "bark_ml/observers/static_observer.hpp"
+#include "gtest/gtest.h"
+
+using namespace bark::models::dynamic;
+using namespace bark::models::behavior;
+using namespace bark::models::execution;
+using namespace bark::world::map;
+using namespace bark_ml::observers;
+
+using bark::commons::SetterParams;
+using bark::commons::transformation::FrenetPosition;
+using bark::commons::transformation::FrenetStateDifference;
+using bark::geometry::Model3D;
+using bark::geometry::Point2d;
+using bark::geometry::Polygon;
+using bark::geometry::Pose;
+using bark::geometry::B_PI;
+using bark::geometry::standard_shapes::CarRectangle;
+using bark::geometry::standard_shapes::GenerateGoalRectangle;
+using bark::world::FrontRearAgents;
+using bark::world::ObservedWorld;
+using bark::world::ObservedWorldPtr;
+using bark::world::World;
+using bark::world::WorldPtr;
+using bark::world::goal_definition::GoalDefinitionPolygon;
+using bark::world::objects::Agent;
+using bark::world::objects::AgentPtr;
+using bark::world::opendrive::OpenDriveMapPtr;
+using bark::world::tests::MakeXodrMapOneRoadTwoLanes;
+
+using StateDefinition::MIN_STATE_SIZE;
+
+const double left_lane_y = -1.75;
+const double right_lane_y = -5.25;
+
+WorldPtr create_world() {
+    auto params = std::make_shared<SetterParams>();
+    WorldPtr world(new World(params));
+    OpenDriveMapPtr open_drive_map = MakeXodrMapOneRoadTwoLanes();
+    MapInterfacePtr map_interface = std::make_shared<MapInterface>();
+    map_interface->interface_from_opendrive(open_drive_map);
+    world->SetMap(map_interface);
+    return world;
+}
+
+AgentId create_agent(WorldPtr world,
+            double x, double y, double v, double theta) {
+    auto params = std::make_shared<SetterParams>();
+    // Goal Definition
+    Polygon polygon = GenerateGoalRectangle(6,3);
+    std::shared_ptr<Polygon> goal_polygon(
+        std::dynamic_pointer_cast<Polygon>(polygon.Translate(Point2d(50, -2))));
+    auto goal_ptr = std::make_shared<GoalDefinitionPolygon>(*goal_polygon);
+
+    // Setting Up Agents (one in front of another)
+    ExecutionModelPtr exec_model(new ExecutionModelInterpolate(params));
+    DynamicModelPtr dyn_model(new SingleTrackModel(params));
+    BehaviorModelPtr beh_model(new BehaviorConstantAcceleration(params));
+    Polygon car_polygon = CarRectangle();
+
+    State init_state(static_cast<int>(MIN_STATE_SIZE));
+    init_state << 0.0, x, y, theta, v;
+    AgentPtr agent(new Agent(init_state, beh_model, dyn_model, exec_model,
+                                car_polygon, params, goal_ptr, world->GetMap(),
+                                Model3D())); 
+    world->AddAgent(agent);
+    return agent->GetAgentId();
+}
+
+
+TEST(nearest_observer, agent_front_same_lane ) { 
+    auto params = std::make_shared<SetterParams>();
+    params->SetInt("ML::NearestObserver::NNearestAgents", 2);
+    params->SetReal("ML::NearestObserver::MinVelLon", -10.0);
+    params->SetReal("ML::NearestObserver::MaxVelLon", 10.0);
+    params->SetReal("ML::NearestObserver::MinVelLat", -10.0);
+    params->SetReal("ML::NearestObserver::MaxVelLat", 10.0);
+    params->SetReal("ML::NearestObserver::MaxDist", 5.0);
+    params->SetReal("ML::NearestObserver::MinS", -20.0);
+    params->SetReal("ML::NearestObserver::MaxS", 20.0);
+    params->SetReal("ML::NearestObserver::MinD", -10.0);
+    params->SetReal("ML::NearestObserver::MaxD", 10.0);
+    params->SetReal("ML::NearestObserver::MinTheta", -B_PI);
+    params->SetReal("ML::NearestObserver::MaxTheta", B_PI);
+    NearestObserver nn_observer(params);
+
+    const auto world = create_world();
+    const auto id1 = create_agent(world, 5.0, left_lane_y, 3.5, 0.0);
+    const auto id2 = create_agent(world, 10.0, left_lane_y, 3.6, 0.0);
+    world->UpdateAgentRTree();
+  
+    const ObservedWorld observed_world(world, id1);
+
+    const auto observed_nn_state = nn_observer.Observe(observed_world);
+    EXPECT_EQ(observed_nn_state.cols(), 6+2*4);
+
+    // Ego lon, lat, theta, vlon, vlat, lane id
+    EXPECT_EQ(observed_nn_state(0, 0), 5.0/20.0);
+    EXPECT_EQ(observed_nn_state(0, 1), 0.0);
+    EXPECT_EQ(observed_nn_state(0, 2), 0.0);
+    EXPECT_EQ(observed_nn_state(0, 3), 3.5/10.0);
+    EXPECT_EQ(observed_nn_state(0, 4), 0.0);
+    EXPECT_EQ(observed_nn_state(0, 5), 0.0);
+
+    // Other dellon, dellat, delvlon, delvlat
+    const auto es = observed_world.GetEgoAgent()->GetShape();
+    EXPECT_EQ(observed_nn_state(0, 6), (10.0 - 5.0 - es.front_dist_ - es.rear_dist_)/20.0);
+    EXPECT_EQ(observed_nn_state(0, 7), 0.0);
+    EXPECT_NEAR(observed_nn_state(0, 8), 0.1/10.0, 0.0001);
+    EXPECT_EQ(observed_nn_state(0, 9), 0.0);
+
+    // Only one other agent in world, rest positions zeroed
+    EXPECT_EQ(observed_nn_state(0, 10), 0.0);
+    EXPECT_EQ(observed_nn_state(0, 11), 0.0);
+    EXPECT_EQ(observed_nn_state(0, 12), 0.0);
+    EXPECT_EQ(observed_nn_state(0, 13), 0.0);
+}
+
+TEST(nearest_observer, agent_front_other_lane_offsets) { 
+    auto params = std::make_shared<SetterParams>();
+    params->SetInt("ML::NearestObserver::NNearestAgents", 4);
+    params->SetReal("ML::NearestObserver::MinVelLon", -10.0);
+    params->SetReal("ML::NearestObserver::MaxVelLon", 10.0);
+    params->SetReal("ML::NearestObserver::MinVelLat", -10.0);
+    params->SetReal("ML::NearestObserver::MaxVelLat", 10.0);
+    params->SetReal("ML::NearestObserver::MaxDist", 50.0);
+    params->SetReal("ML::NearestObserver::MinS", -50.0);
+    params->SetReal("ML::NearestObserver::MaxS", 50.0);
+    params->SetReal("ML::NearestObserver::MinD", -15.0);
+    params->SetReal("ML::NearestObserver::MaxD", 15.0);
+    params->SetReal("ML::NearestObserver::MinTheta", -B_PI);
+    params->SetReal("ML::NearestObserver::MaxTheta", B_PI);
+    NearestObserver nn_observer(params);
+
+    const auto world = create_world();
+    const auto id1 = create_agent(world, 20.0, right_lane_y-0.5, 3.8, 0.4);
+    const auto id2 = create_agent(world, 26.0, left_lane_y+1.0, 4.6, -0.6);
+    const auto id3 = create_agent(world, 35.0, left_lane_y+1.0, 4.6, -0.6);
+    const auto id4 = create_agent(world, 100.0, 0, 0, 0); // outside distance limit
+    world->UpdateAgentRTree();
+  
+    const ObservedWorld observed_world(world, id1);
+
+    const auto observed_nn_state = nn_observer.Observe(observed_world);
+    EXPECT_EQ(observed_nn_state.cols(), 6+4*4);
+
+    // Ego lon, lat, theta, vlon, vlat, lane id
+    EXPECT_EQ(observed_nn_state(0, 0), 20.0/50.0);
+    EXPECT_EQ(observed_nn_state(0, 1), -0.5/15.0);
+    EXPECT_NEAR(observed_nn_state(0, 2), -0.4/B_PI, 0.0001);
+    EXPECT_NEAR(observed_nn_state(0, 3), 3.8*cos(0.4)/10.0, 0.0001);
+    EXPECT_NEAR(observed_nn_state(0, 4), 3.8*sin(0.4)/10.0, 0.0001);
+    EXPECT_EQ(observed_nn_state(0, 5), 1.0);
+
+    // Other agent 1 dellon, dellat, delvlon, delvlat
+    const auto es = observed_world.GetEgoAgent()->GetShape();
+    EXPECT_NEAR(observed_nn_state(0, 6), (26.0 - 20.0 - es.front_dist_*cos(0.4) - es.rear_dist_*cos(0.6) -
+                                             es.left_dist_*sin(0.4) - es.left_dist_*sin(0.6)  )/50.0, 0.0001);
+    EXPECT_NEAR(observed_nn_state(0, 7), (left_lane_y +  1.0 - right_lane_y + 0.5 - es.rear_dist_*sin(0.4)-
+                                             es.front_dist_*sin(0.6) - es.left_dist_*cos(0.4) -
+                                             es.left_dist_*cos(0.6)    )/15.0, 0.0001);
+    EXPECT_NEAR(observed_nn_state(0, 8), -(3.8*cos(0.4) - 4.6*cos(0.6) )/10.0, 0.0001);
+    EXPECT_NEAR(observed_nn_state(0, 9), -(3.8*sin(0.4) + 4.6*sin(0.6) )/10.0, 0.0001);
+
+    // Other agent 2 dellon, dellat, delvlon, delvlat
+    EXPECT_NEAR(observed_nn_state(0, 10), (35.0 - 20.0 - es.front_dist_*cos(0.4) - es.rear_dist_*cos(0.6) -
+                                             es.left_dist_*sin(0.4) - es.left_dist_*sin(0.6)  )/50.0, 0.0001);
+    EXPECT_NEAR(observed_nn_state(0, 11), (left_lane_y +  1.0 - right_lane_y + 0.5 - es.rear_dist_*sin(0.4)-
+                                             es.front_dist_*sin(0.6) - es.left_dist_*cos(0.4) -
+                                             es.left_dist_*cos(0.6)    )/15.0, 0.0001);
+    EXPECT_NEAR(observed_nn_state(0, 12), -(3.8*cos(0.4) - 4.6*cos(0.6) )/10.0, 0.0001);
+    EXPECT_NEAR(observed_nn_state(0, 13), -(3.8*sin(0.4) + 4.6*sin(0.6) )/10.0, 0.0001);
+
+
+    // Only nearest agents considered, third agent out of max dist -> positions zeroed
+    EXPECT_EQ(observed_nn_state(0, 14), 0.0);
+    EXPECT_EQ(observed_nn_state(0, 15), 0.0);
+    EXPECT_EQ(observed_nn_state(0, 16), 0.0);
+    EXPECT_EQ(observed_nn_state(0, 17), 0.0);
+}
+
diff --git a/bark_ml/tests/observer_test.cc b/bark_ml/tests/observer_static_test.cc
similarity index 98%
rename from bark_ml/tests/observer_test.cc
rename to bark_ml/tests/observer_static_test.cc
index 39c0db75..9c1f2173 100644
--- a/bark_ml/tests/observer_test.cc
+++ b/bark_ml/tests/observer_static_test.cc
@@ -12,6 +12,7 @@
 #include "bark/geometry/standard_shapes.hpp"
 #include "bark/models/behavior/constant_acceleration/constant_acceleration.hpp"
 #include "bark/models/behavior/motion_primitives/continuous_actions.hpp"
+#include "bark/models/dynamic/single_track.hpp"
 #include "bark/models/dynamic/single_track_steering_rate.hpp"
 #include "bark/models/execution/interpolation/interpolate.hpp"
 #include "bark/world/evaluation/evaluator_collision_agents.hpp"
@@ -107,6 +108,7 @@ TEST(static_observer, agents_exist ) {
     params->SetReal("ML::StaticObserver::MaxD", 10.0);
     params->SetReal("ML::StaticObserver::MinTheta", -B_PI);
     params->SetReal("ML::StaticObserver::MaxTheta", B_PI);
+    params->SetBool("ML::StaticObserver::CheckInRoadCorridor", false);
     StaticObserver nn_observer(params);
 
     const auto world = create_world();
diff --git a/bark_ml/tests/py_graph_nets_tests.py b/bark_ml/tests/py_graph_nets_tests.py
deleted file mode 100644
index a1d8580d..00000000
--- a/bark_ml/tests/py_graph_nets_tests.py
+++ /dev/null
@@ -1,94 +0,0 @@
-# Copyright (c) 2020 fortiss GmbH
-#
-# Authors: Patrick Hart, Julian Bernhard, Klemens Esterle, and
-# Tobias Kessler
-#
-# This software is released under the MIT License.
-# https://opensource.org/licenses/MIT
-
-
-import unittest
-import numpy as np
-from graph_nets import utils_tf
-from bark.runtime.commons.parameters import ParameterServer
-
-from bark_ml.observers.graph_observer import GraphObserver
-
-
-class PyGraphNetsTests(unittest.TestCase):
-  def test_gnn(self):
-    # Node features for graph 0.
-    nodes_0 = [[10.1, 20., 30.],  # Node 0
-              [11., 21., 31.],  # Node 1
-              [12., 22., 32.],  # Node 2
-              [13., 23., 33.],  # Node 3
-              [14., 24., 34.]]  # Node 4
-
-    # Edge features for graph 0.
-    edges_0 = [[100., 200.],  # Edge 0
-              [101.2, 201.],  # Edge 1
-              [102., 202.],  # Edge 2
-              [103., 203.],  # Edge 3
-              [104., 204.],  # Edge 4
-              [105., 205.]]  # Edge 5
-
-    # The sender and receiver nodes associated with each edge for graph 0.
-    senders_0 = [0,  # Index of the sender node for edge 0
-                1,  # Index of the sender node for edge 1
-                1,  # Index of the sender node for edge 2
-                2,  # Index of the sender node for edge 3
-                2,  # Index of the sender node for edge 4
-                3]  # Index of the sender node for edge 5
-    receivers_0 = [1,  # Index of the receiver node for edge 0
-                  2,  # Index of the receiver node for edge 1
-                  3,  # Index of the receiver node for edge 2
-                  0,  # Index of the receiver node for edge 3
-                  3,  # Index of the receiver node for edge 4
-                  4]  # Index of the receiver node for edge 5
-
-    data_dict_0 = {
-      "globals": [],
-      "nodes": nodes_0,
-      "edges": edges_0,
-      "senders": senders_0,
-      "receivers": receivers_0
-    }
-
-    input_graph = utils_tf.data_dicts_to_graphs_tuple(
-      [data_dict_0, data_dict_0])
-
-    num_nodes = len(nodes_0)
-    num_features = 3
-    num_edge_features = len(edges_0)
-    graph_dims = (num_nodes, num_features, num_edge_features)
-
-    # 6 edges x 2
-    # 5 nodes x 3
-    # adj matrix 5x5
-    obs = np.zeros(shape=(1, 52))
-    # NOTE: use dense
-    params = ParameterServer()
-    graph_observer = GraphObserver(params)
-    graph_observer.feature_len = 2
-    graph_observer.edge_feature_len = 3
-
-    _, _, _ = graph_observer.graph(obs, graph_dims=graph_dims)
-
-    print(input_graph)
-    # gnn = MLPGraphNetwork()
-    # input_graph_1 = GraphsTuple(
-    #   nodes=tf.convert_to_tensor(nodes_0, dtype=tf.float32),
-    #   edges=tf.convert_to_tensor(edges_0, dtype=tf.float32),
-    #   globals=tf.cast(tf.tile([[0]], [1, 1]), tf.float32),
-    #   receivers=tf.convert_to_tensor(receivers_0, dtype=tf.int32),
-    #   senders=tf.convert_to_tensor(senders_0, dtype=tf.int32),
-    #   n_node=tf.tile([5], [1]),
-    #   n_edge=tf.tile([6], [1]))
-
-    # tf.print(input_graph_1)
-    # result = gnn(input_graph_1)
-    # tf.print(result)
-
-
-if __name__ == '__main__':
-  unittest.main()
\ No newline at end of file
diff --git a/bark_ml/tests/py_library_tfa_tests.py b/bark_ml/tests/py_library_tfa_tests.py
index 3c12c161..0ddc728d 100644
--- a/bark_ml/tests/py_library_tfa_tests.py
+++ b/bark_ml/tests/py_library_tfa_tests.py
@@ -4,6 +4,11 @@
 #
 # This work is licensed under the terms of the MIT license.
 # For a copy, see <https://opensource.org/licenses/MIT>.
+try:
+    import debug_settings
+except:
+    pass
+
 
 import unittest
 import numpy as np
diff --git a/bark_ml/tests/py_library_torch_agents_tests.py b/bark_ml/tests/py_library_torch_agents_tests.py
deleted file mode 100644
index 750b5f5f..00000000
--- a/bark_ml/tests/py_library_torch_agents_tests.py
+++ /dev/null
@@ -1,187 +0,0 @@
-# Copyright (c) 2020 fortiss GmbH
-#
-# Authors: Patrick Hart
-#
-# This work is licensed under the terms of the MIT license.
-# For a copy, see <https://opensource.org/licenses/MIT>.
-
-import unittest
-import numpy as np
-import os
-import gym
-import matplotlib
-import time
-
-# BARK imports
-from bark.runtime.commons.parameters import ParameterServer
-
-# BARK-ML imports
-from bark_ml.environments.blueprints import \
-  DiscreteHighwayBlueprint
-from bark_ml.environments.single_agent_runtime import SingleAgentRuntime
-import bark_ml.environments.gym  # pylint: disable=unused-import
-from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.model_wrapper \
- import pytorch_script_wrapper
-from bark_ml.library_wrappers.lib_fqf_iqn_qrdqn.agent import IQNAgent, FQFAgent, QRDQNAgent
-
-
-class PyLibraryWrappersPyTorchAgentTests(unittest.TestCase):
-  """TorchAgentTests tests."""
-
-  # make sure the agent works
-  def test_agent_wrapping(self):
-    params = ParameterServer()
-    env = gym.make("highway-v1", params=params)
-    env.reset()
-    agent = IQNAgent(env=env, test_env=env, params=params)
-    agent = FQFAgent(env=env, test_env=env, params=params)
-    agent = QRDQNAgent(env=env, test_env=env, params=params)
-
-  # assign as behavior model (to check if trained agent can be used)
-  def test_behavior_wrapping(self):
-    # create scenario
-    params = ParameterServer()
-    bp = DiscreteHighwayBlueprint(params, num_scenarios=10, random_seed=0)
-    env = SingleAgentRuntime(blueprint=bp, render=False)
-    #env = gym.make("highway-v1", params=params)
-    ml_behaviors = []
-    ml_behaviors.append(IQNAgent(env=env, test_env=env, params=params))
-    ml_behaviors.append(FQFAgent(env=env, test_env=env, params=params))
-    ml_behaviors.append(QRDQNAgent(env=env, test_env=env, params=params))
-
-    for ml_behavior in ml_behaviors:
-      # set agent
-      env.ml_behavior = ml_behavior
-      env.reset()
-      done = False
-      while done is False:
-        action = np.random.randint(low=0, high=env.action_space.n)
-        observed_next_state, reward, done, info = env.step(action)
-        print(
-            f"Observed state: {observed_next_state}, Reward: {reward}, Done: {done}"
-        )
-
-      # action is set externally
-      ml_behavior._set_action_externally = True
-      agent_id = list(env._world.agents.keys())[0]
-      observed_world = env._world.Observe([agent_id])[0]
-
-      # do a random action and plan trajectory
-      action = np.random.randint(low=1, high=env.action_space.n)
-      ml_behavior.ActionToBehavior(action)
-      a = ml_behavior.Plan(0.2, observed_world)
-
-      # sample another different random action
-      another_action = action
-      while another_action == action:
-        another_action = np.random.randint(low=1, high=env.action_space.n)
-
-      # plan trajectory for the another action
-      ml_behavior.ActionToBehavior(another_action)
-      b = ml_behavior.Plan(0.2, observed_world)
-
-      # the trajectory generated by two different actions shoould be different
-      self.assertEqual(np.any(np.not_equal(a, b)), True)
-
-      # action will be calculated within the Plan(..) fct.
-      ml_behavior._set_action_externally = False
-      a = ml_behavior.Plan(0.2, observed_world)
-      b = ml_behavior.Plan(0.2, observed_world)
-
-      # same trajectory for same state
-      np.testing.assert_array_equal(a, b)
-
-  def test_agents(self):
-    params = ParameterServer()
-    params["ML"]["BaseAgent"]["NumSteps"] = 2
-    params["ML"]["BaseAgent"]["MaxEpisodeSteps"] = 2
-
-    bp = DiscreteHighwayBlueprint(params, num_scenarios=10, random_seed=0)
-    env = SingleAgentRuntime(blueprint=bp, render=False)
-
-    # IQN Agent
-    iqn_agent = IQNAgent(env=env, test_env=env, params=params)
-    env.ml_behavior = iqn_agent
-    self.assertEqual(env.ml_behavior.set_action_externally, False)
-    iqn_agent.run()
-    self.assertEqual(env.ml_behavior.set_action_externally, True)
-
-    # FQF Agent
-    fqf_agent = FQFAgent(env=env, test_env=env, params=params)
-    env.ml_behavior = fqf_agent
-    self.assertEqual(env.ml_behavior.set_action_externally, False)
-    fqf_agent.run()
-    self.assertEqual(env.ml_behavior.set_action_externally, True)
-
-    # QRDQN Agent
-    qrdqn_agent = QRDQNAgent(env=env, test_env=env, params=params)
-    env.ml_behavior = qrdqn_agent
-    self.assertEqual(env.ml_behavior.set_action_externally, False)
-    qrdqn_agent.run()
-    self.assertEqual(env.ml_behavior.set_action_externally, True)
-
-  def test_model_loader(self):
-    # env using default params
-    env = gym.make("highway-v1")
-
-    networks = ["iqn", "fqf", "qrdqn"]
-
-    action_space_size = env.action_space.n
-    state_space_size = env.observation_space.shape[0]
-
-    # a sample random state [0-1] to evaluate actions
-    random_state = np.random.rand(state_space_size).tolist()
-
-    # test all networks
-    for network in networks:
-      # Do inference using C++ wrapped model
-      model = pytorch_script_wrapper.ModelLoader(
-          os.path.join(
-              os.path.dirname(__file__),
-              "lib_fqf_iqn_qrdqn_test_data/{}/online_net_script.pt"
-              .format(network)), action_space_size, state_space_size)
-      model.LoadModel()
-
-      num_iters = 1000  # Number of times to repeat experiment to calcualte runtime
-
-      # Time num_iters iterations for inference using C++ model
-      start = time.time()
-      for _ in range(num_iters):
-        actions_cpp = model.Inference(random_state)
-      end = time.time()
-      time_cpp = end - start  # todo - how to analyze python vs cpp test time in tests?
-
-      # Load and perform inference using python model
-      if network == "iqn":
-        agent = IQNAgent(env=env, test_env=env, params=ParameterServer())
-
-      elif network == "fqf":
-        agent = FQFAgent(env=env, test_env=env, params=ParameterServer())
-
-      elif network == "qrdqn":
-        agent = QRDQNAgent(env=env, test_env=env, params=ParameterServer())
-
-      agent.load_models(
-          os.path.join(
-              os.path.dirname(__file__),
-              "lib_fqf_iqn_qrdqn_test_data",
-              network))
-
-      # Time num_iters iterations for inference using python model
-      start = time.time()
-      for _ in range(num_iters):
-        actions_py = agent.calculate_actions(random_state)
-
-      end = time.time()
-      time_py = end - start
-
-      # assert that Python and Cpp models are close enough to 6 decimal places
-      np.testing.assert_array_almost_equal(
-          actions_py.flatten().numpy(),
-          np.asarray(actions_cpp),
-          decimal=6,
-          err_msg="C++ and python models don't match")
-
-
-if __name__ == '__main__':
-  unittest.main()
diff --git a/test_demo_collected b/test_demo_collected
new file mode 100644
index 00000000..0dda67a3
Binary files /dev/null and b/test_demo_collected differ