diff --git a/examples/shields/09_dtcontrol.py b/examples/shields/09_pre_shield_decision_tree.py
similarity index 65%
rename from examples/shields/09_dtcontrol.py
rename to examples/shields/09_pre_shield_decision_tree.py
index b27a9fe..8a9b158 100644
--- a/examples/shields/09_dtcontrol.py
+++ b/examples/shields/09_pre_shield_decision_tree.py
@@ -8,7 +8,7 @@ import stormpy.shields
 import stormpy.examples
 import stormpy.examples.files
 
-from stormpy.dtcontrol import export_decision_tree
+from stormpy.decision_tree import create_decision_tree
 
 def export_shield_as_dot():
     path = stormpy.examples.files.prism_mdp_lava_simple
@@ -24,17 +24,20 @@ def export_shield_as_dot():
     options.set_build_with_choice_origins(True)
     model = stormpy.build_sparse_model_with_options(program, options)
 
-    result = stormpy.model_checking(model, formulas[0], extract_scheduler=True) #, shielding_expression=shield_specification)
-    
+    result = stormpy.model_checking(model, formulas[0], extract_scheduler=True) 
+
     assert result.has_shield
 
     shield = result.shield
-    stormpy.shields.export_shieldDouble(model, shield, "preshield.storm.json")
-
-
-    export_decision_tree(result.shield)
-
-
+    filename = "preshield.storm.json"
+    filename2 = "preshield.shield"
+    stormpy.shields.export_shieldDouble(model, shield, filename)
+    stormpy.shields.export_shieldDouble(model, shield, filename2)
+    
+    output_folder = "pre_trees"
+    name = 'pre_my_output'
+    suite = create_decision_tree(filename, name=name , output_folder=output_folder, export_pdf=True)
+    suite.display_html()
 
 if __name__ == '__main__':
     export_shield_as_dot()
\ No newline at end of file
diff --git a/examples/shields/10_post_shield_decision_tree.py b/examples/shields/10_post_shield_decision_tree.py
new file mode 100644
index 0000000..0bf3100
--- /dev/null
+++ b/examples/shields/10_post_shield_decision_tree.py
@@ -0,0 +1,43 @@
+import stormpy
+import stormpy.core
+import stormpy.simulator
+
+
+import stormpy.shields
+
+import stormpy.examples
+import stormpy.examples.files
+
+from stormpy.decision_tree import create_decision_tree
+
+def export_shield_as_dot():
+    path = stormpy.examples.files.prism_mdp_lava_simple
+    formula_str = "<post, PostSafety, gamma=1> Pmax=? [G !\"AgentIsInLavaAndNotDone\"]"
+
+    program = stormpy.parse_prism_program(path)
+    formulas = stormpy.parse_properties_for_prism_program(formula_str, program)
+
+    options = stormpy.BuilderOptions([p.raw_formula for p in formulas])
+    options.set_build_state_valuations(True)
+    options.set_build_choice_labels(True)
+    options.set_build_all_labels()
+    options.set_build_with_choice_origins(True)
+    model = stormpy.build_sparse_model_with_options(program, options)
+
+    result = stormpy.model_checking(model, formulas[0], extract_scheduler=True) 
+
+    assert result.has_shield
+
+    shield = result.shield
+    filename = "postshield.storm.json"
+    filename2 = "postshield.shield"
+    stormpy.shields.export_shieldDouble(model, shield, filename)
+    stormpy.shields.export_shieldDouble(model, shield, filename2)
+    
+    output_folder = "post_trees"
+    name = 'post_my_output'
+    suite = create_decision_tree(filename, name=name , output_folder=output_folder, export_pdf=True)
+    suite.display_html()
+
+if __name__ == '__main__':
+    export_shield_as_dot()
\ No newline at end of file
diff --git a/lib/stormpy/decision_tree.py b/lib/stormpy/decision_tree.py
new file mode 100644
index 0000000..281abed
--- /dev/null
+++ b/lib/stormpy/decision_tree.py
@@ -0,0 +1,43 @@
+from sklearn.linear_model import LogisticRegression
+from dtcontrol.benchmark_suite import BenchmarkSuite
+from dtcontrol.decision_tree.decision_tree import DecisionTree
+from dtcontrol.decision_tree.determinization.max_freq_determinizer import MaxFreqDeterminizer
+from dtcontrol.decision_tree.impurity.entropy import Entropy
+from dtcontrol.decision_tree.impurity.multi_label_entropy import MultiLabelEntropy
+from dtcontrol.decision_tree.splitting.axis_aligned import AxisAlignedSplittingStrategy
+from dtcontrol.decision_tree.splitting.linear_classifier import LinearClassifierSplittingStrategy
+
+import pydot
+
+def create_decision_tree(filename, name, output_folder, 
+                         timeout=60*60*2,
+                         benchmark_file='benchmark',
+                         save_folder='saved_classifiers',
+                         export_pdf=False,
+                         classifiers=None):
+    
+    suite = BenchmarkSuite(timeout=timeout,
+                            save_folder=save_folder,
+                            output_folder=output_folder,
+                            benchmark_file=benchmark_file,
+                            rerun=True)
+
+    suite.add_datasets([filename])  
+
+
+    if classifiers is None:        
+        aa = AxisAlignedSplittingStrategy()
+        aa.priority = 1      
+
+        classifiers = [DecisionTree([aa], Entropy(), name)]
+
+    suite.benchmark(classifiers)
+    if export_pdf:
+        for dataset in suite.datasets:
+            for classifier in classifiers:
+                filename = suite.get_filename(output_folder, dataset=dataset , classifier=classifier, extension='.dot')
+                (graph,) = pydot.graph_from_dot_file(filename)
+                graph.write_pdf(F'{name}.pdf')
+
+
+    return suite
diff --git a/lib/stormpy/dtcontrol.py b/lib/stormpy/dtcontrol.py
deleted file mode 100644
index 1f28158..0000000
--- a/lib/stormpy/dtcontrol.py
+++ /dev/null
@@ -1,17 +0,0 @@
-import json
-import logging
-from os.path import splitext, exists
-from sklearn.linear_model import LogisticRegression
-import numpy as np
-
-from dtcontrol.dataset.dataset_loader import DatasetLoader
-from dtcontrol.decision_tree.decision_tree import DecisionTree
-from dtcontrol.decision_tree.impurity.entropy import Entropy
-
-
-from dtcontrol.decision_tree.splitting.axis_aligned import AxisAlignedSplittingStrategy
-from dtcontrol.decision_tree.splitting.linear_classifier import LinearClassifierSplittingStrategy
-
-
-def export_decision_tree(filename):
-    pass