diff --git a/pomdp_problems/__init__.py b/pomdp_problems/__init__.py
deleted file mode 100644
index 58259f9c..00000000
--- a/pomdp_problems/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from pomdp_problems import tiger
diff --git a/pomdp_problems/light_dark/__init__.py b/pomdp_problems/light_dark/__init__.py
deleted file mode 100644
index 18c82386..00000000
--- a/pomdp_problems/light_dark/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-from pomdp_problems.light_dark.env.env import LightDarkEnvironment
-from pomdp_problems.light_dark.env.visual import LightDarkViz
-from pomdp_problems.light_dark.domain.state import *
-from pomdp_problems.light_dark.domain.action import *
-from pomdp_problems.light_dark.domain.observation import *
-from pomdp_problems.light_dark.models.transition_model import *
-from pomdp_problems.light_dark.models.observation_model import *
diff --git a/pomdp_problems/tiger/__init__.py b/pomdp_problems/tiger/__init__.py
deleted file mode 100644
index 48a81b64..00000000
--- a/pomdp_problems/tiger/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from .tiger_problem import TigerProblem, test_planner
diff --git a/pomdp_py/__init__.py b/pomdp_py/__init__.py
index cf75b214..b2ce26e2 100644
--- a/pomdp_py/__init__.py
+++ b/pomdp_py/__init__.py
@@ -22,7 +22,6 @@
 from pomdp_py.algorithms.po_uct import POUCT, QNode, VNode, RootVNode,\
     RolloutPolicy, RandomRollout, ActionPrior
 from pomdp_py.algorithms.bsp.blqr import BLQR
-from pomdp_py import visual
 
 # Templates & Utilities
 from pomdp_py.utils.templates import *
diff --git a/pomdp_py/__main__.py b/pomdp_py/__main__.py
index c38f04e9..832dd93d 100644
--- a/pomdp_py/__main__.py
+++ b/pomdp_py/__main__.py
@@ -11,7 +11,7 @@
 def parse_args():
     parser = argparse.ArgumentParser(description="pomdp_py CLI")
     parser.add_argument("-r", "--run", type=str,
-                        help="run a pomdp under pomdp_problems."
+                        help="run a pomdp under pomdp_py.problems."
                         "Available options: {}".format(available_problems))
     args = parser.parse_args()
     return parser, args
@@ -20,23 +20,23 @@ def parse_args():
     parser, args = parse_args()
     if args.run:
         if args.run.lower() == "tiger":
-            from pomdp_problems.tiger.tiger_problem import main
+            from pomdp_py.problems.tiger.tiger_problem import main
             main()
 
         elif args.run.lower() == "rocksample":
-            from pomdp_problems.rocksample.rocksample_problem import main
+            from pomdp_py.problems.rocksample.rocksample_problem import main
             main()
 
         elif args.run.lower() == "mos":
-            from pomdp_problems.multi_object_search.problem import unittest
+            from pomdp_py.problems.multi_object_search.problem import unittest
             unittest()
 
         elif args.run.lower() == "tag":
-            from pomdp_problems.tag.problem import main
+            from pomdp_py.problems.tag.problem import main
             main()
 
         elif args.run.lower() == "load_unload":
-            from pomdp_problems.load_unload.load_unload import main
+            from pomdp_py.problems.load_unload.load_unload import main
             main()
 
         else:
diff --git a/pomdp_py/problems/__init__.py b/pomdp_py/problems/__init__.py
new file mode 100644
index 00000000..ed890cb6
--- /dev/null
+++ b/pomdp_py/problems/__init__.py
@@ -0,0 +1 @@
+from pomdp_py.problems import tiger
diff --git a/pomdp_problems/lasertag/README.rst b/pomdp_py/problems/lasertag/README.rst
similarity index 100%
rename from pomdp_problems/lasertag/README.rst
rename to pomdp_py/problems/lasertag/README.rst
diff --git a/pomdp_problems/lasertag/__init__.py b/pomdp_py/problems/lasertag/__init__.py
similarity index 100%
rename from pomdp_problems/lasertag/__init__.py
rename to pomdp_py/problems/lasertag/__init__.py
diff --git a/pomdp_problems/light_dark/README.rst b/pomdp_py/problems/light_dark/README.rst
similarity index 100%
rename from pomdp_problems/light_dark/README.rst
rename to pomdp_py/problems/light_dark/README.rst
diff --git a/pomdp_py/problems/light_dark/__init__.py b/pomdp_py/problems/light_dark/__init__.py
new file mode 100644
index 00000000..528a4a5f
--- /dev/null
+++ b/pomdp_py/problems/light_dark/__init__.py
@@ -0,0 +1,7 @@
+from pomdp_py.problems.light_dark.env.env import LightDarkEnvironment
+from pomdp_py.problems.light_dark.env.visual import LightDarkViz
+from pomdp_py.problems.light_dark.domain.state import *
+from pomdp_py.problems.light_dark.domain.action import *
+from pomdp_py.problems.light_dark.domain.observation import *
+from pomdp_py.problems.light_dark.models.transition_model import *
+from pomdp_py.problems.light_dark.models.observation_model import *
diff --git a/pomdp_problems/light_dark/agent/__init__.py b/pomdp_py/problems/light_dark/agent/__init__.py
similarity index 100%
rename from pomdp_problems/light_dark/agent/__init__.py
rename to pomdp_py/problems/light_dark/agent/__init__.py
diff --git a/pomdp_problems/light_dark/agent/belief.py b/pomdp_py/problems/light_dark/agent/belief.py
similarity index 100%
rename from pomdp_problems/light_dark/agent/belief.py
rename to pomdp_py/problems/light_dark/agent/belief.py
diff --git a/pomdp_problems/light_dark/domain/__init__.py b/pomdp_py/problems/light_dark/domain/__init__.py
similarity index 100%
rename from pomdp_problems/light_dark/domain/__init__.py
rename to pomdp_py/problems/light_dark/domain/__init__.py
diff --git a/pomdp_problems/light_dark/domain/action.py b/pomdp_py/problems/light_dark/domain/action.py
similarity index 100%
rename from pomdp_problems/light_dark/domain/action.py
rename to pomdp_py/problems/light_dark/domain/action.py
diff --git a/pomdp_problems/light_dark/domain/observation.py b/pomdp_py/problems/light_dark/domain/observation.py
similarity index 100%
rename from pomdp_problems/light_dark/domain/observation.py
rename to pomdp_py/problems/light_dark/domain/observation.py
diff --git a/pomdp_problems/light_dark/domain/state.py b/pomdp_py/problems/light_dark/domain/state.py
similarity index 100%
rename from pomdp_problems/light_dark/domain/state.py
rename to pomdp_py/problems/light_dark/domain/state.py
diff --git a/pomdp_problems/light_dark/env/__init__.py b/pomdp_py/problems/light_dark/env/__init__.py
similarity index 100%
rename from pomdp_problems/light_dark/env/__init__.py
rename to pomdp_py/problems/light_dark/env/__init__.py
diff --git a/pomdp_problems/light_dark/env/env.py b/pomdp_py/problems/light_dark/env/env.py
similarity index 96%
rename from pomdp_problems/light_dark/env/env.py
rename to pomdp_py/problems/light_dark/env/env.py
index c0296f69..da4930dd 100644
--- a/pomdp_problems/light_dark/env/env.py
+++ b/pomdp_py/problems/light_dark/env/env.py
@@ -4,11 +4,11 @@
 """
 
 import pomdp_py
-import pomdp_problems.light_dark as ld
+import pomdp_py.problems.light_dark as ld
 import numpy as np
 
 class LightDarkEnvironment(pomdp_py.Environment):
-    
+
     def __init__(self,
                  init_state,
                  light,
diff --git a/pomdp_py/utils/plotting.py b/pomdp_py/problems/light_dark/env/plotting.py
similarity index 95%
rename from pomdp_py/utils/plotting.py
rename to pomdp_py/problems/light_dark/env/plotting.py
index 9cf3ff89..f347dac0 100644
--- a/pomdp_py/utils/plotting.py
+++ b/pomdp_py/problems/light_dark/env/plotting.py
@@ -14,10 +14,6 @@ def plot_points(xvals, yvals, color=None,
     plt.ylabel(ylabel)
     plt.legend(loc=loc)
 
-def save_plot(path, bbox_inches='tight'):
-    plt.savefig(path, bbox_inches=bbox_inches)
-    plt.close()
-
 
 # Plot polygons with colors
 def plot_polygons(verts, colors, ax=None, edgecolor=None):
diff --git a/pomdp_problems/light_dark/env/visual.py b/pomdp_py/problems/light_dark/env/visual.py
similarity index 99%
rename from pomdp_problems/light_dark/env/visual.py
rename to pomdp_py/problems/light_dark/env/visual.py
index 9f02a11b..0e497a5d 100644
--- a/pomdp_problems/light_dark/env/visual.py
+++ b/pomdp_py/problems/light_dark/env/visual.py
@@ -1,7 +1,7 @@
 """Plot the light dark environment"""
 import matplotlib.pyplot as plt
 from matplotlib.collections import PolyCollection
-import pomdp_problems.light_dark as ld
+import pomdp_py.problems.light_dark as ld
 from pomdp_py.utils import plotting, colors
 from pomdp_py.utils.misc import remap
 
diff --git a/pomdp_problems/light_dark/models/__init__.py b/pomdp_py/problems/light_dark/models/__init__.py
similarity index 100%
rename from pomdp_problems/light_dark/models/__init__.py
rename to pomdp_py/problems/light_dark/models/__init__.py
diff --git a/pomdp_problems/light_dark/models/observation_model.py b/pomdp_py/problems/light_dark/models/observation_model.py
similarity index 100%
rename from pomdp_problems/light_dark/models/observation_model.py
rename to pomdp_py/problems/light_dark/models/observation_model.py
diff --git a/pomdp_problems/light_dark/models/transition_model.py b/pomdp_py/problems/light_dark/models/transition_model.py
similarity index 100%
rename from pomdp_problems/light_dark/models/transition_model.py
rename to pomdp_py/problems/light_dark/models/transition_model.py
diff --git a/pomdp_problems/light_dark/problem_discrete.py b/pomdp_py/problems/light_dark/problem_discrete.py
similarity index 100%
rename from pomdp_problems/light_dark/problem_discrete.py
rename to pomdp_py/problems/light_dark/problem_discrete.py
diff --git a/pomdp_problems/light_dark/test.py b/pomdp_py/problems/light_dark/test.py
similarity index 96%
rename from pomdp_problems/light_dark/test.py
rename to pomdp_py/problems/light_dark/test.py
index 20c5dd2f..48137926 100644
--- a/pomdp_problems/light_dark/test.py
+++ b/pomdp_py/problems/light_dark/test.py
@@ -1,4 +1,4 @@
-import pomdp_problems.light_dark as ld
+import pomdp_py.problems.light_dark as ld
 import matplotlib.pyplot as plt
 import pomdp_py
 import numpy as np
@@ -77,7 +77,7 @@ def manual_test(blqr):
         print(b_1[1])
         b_1 = blqr.ekf_update_mlo(b_1, [-1.0, 0.0])
         print(b_1[0])
-        print(b_1[1])    
+        print(b_1[1])
         b_1 = blqr.ekf_update_mlo(b_1, [-1.0, 0.0])
         print(b_1[0])
         print(b_1[1])
@@ -126,7 +126,7 @@ def manual_test(blqr):
         b_traj_light.append(b_tp1)
         b_t = b_tp1
     bu_traj_light = [(b_traj_light[t], np.array(u_traj_light[t])) for t in range(len(u_traj_light))]
-    
+
     # traj not through light
     u_traj_dark = [[-1., -1.], [-1., -1.], [-1., -1.]]
     b_t = b_0
@@ -135,12 +135,12 @@ def manual_test(blqr):
         u_t = u_traj_dark[t]
         b_tp1 = blqr.ekf_update_mlo(b_t, u_t)
         b_traj_dark.append(b_tp1)
-        b_t = b_tp1        
+        b_t = b_tp1
     bu_traj_dark = [(b_traj_dark[t], np.array(u_traj_dark[t])) for t in range(len(u_traj_dark))]
 
 
     total_light = 0
-    total_dark = 0    
+    total_dark = 0
     for i in range(1000):
         cost_light = blqr.segmented_cost_function(bu_traj_light, b_des, [], len(bu_traj_light))
         cost_dark = blqr.segmented_cost_function(bu_traj_dark, b_des, [], len(bu_traj_dark))
@@ -154,7 +154,7 @@ def manual_test(blqr):
     viz = ld.LightDarkViz(env, x_range, y_range, 0.1)
     viz.set_goal(goal_pos)
     viz.set_initial_belief_pos(b_0[0])
-    
+
     # Made up paths
     viz.log_position(tuple(b_0[0]), path=2)
     viz.log_position(tuple(b_0[0]), path=3)
@@ -171,7 +171,7 @@ def manual_test(blqr):
     for b_t, u_t in bu_traj_dark:
         viz.log_position(tuple(b_t[0]), path=4)
         viz.log_position(tuple(sysd_b_dark[-1][0]), path=5)
-        sysd_b_dark.append((func_sysd(sysd_b_dark[-1][0], u_t), 0))        
+        sysd_b_dark.append((func_sysd(sysd_b_dark[-1][0], u_t), 0))
 
     viz.plot(path_colors={2: [(0,0,0), (0,255,0)],
                           3: [(0,0,0), (0,255,255)],
@@ -181,7 +181,7 @@ def manual_test(blqr):
                           3: "-",
                           4: "--",
                           5: "-"})
-    plt.show()    
+    plt.show()
 
 
 
@@ -197,13 +197,13 @@ def test(blqr):
     ############
     b_0 = (np.array([2.0, 2.0]),
            np.array([[5.0, 0.0],
-                     [0.0, 5.0]]))    
+                     [0.0, 5.0]]))
     x_sol = blqr.create_plan(b_0, b_des, u_des,
                              num_segments=num_segments,
                              opt_options={"maxiter": 30},
                              opt_callback=opt_callback,
                              control_bounds=(-0.1, 0.1))
-    with np.printoptions(precision=3, suppress=True):    
+    with np.printoptions(precision=3, suppress=True):
         print("SLSQP solution:")
         print(x_sol)
     plan = blqr.interpret_sqp_plan(x_sol, num_segments)
@@ -227,22 +227,22 @@ def test(blqr):
                           1: "-"},
              path_widths={0: 4,
                           1: 1})
-    plt.show()    
+    plt.show()
+
 
-    
 
 
 if __name__ == "__main__":
     test(blqr)
     # bt = (np.array([1.3, 2.0]), np.array([[0.5, 0.0], [0.0, 0.5]]))
     # ut = np.array([1.5, 0.2])
-    
+
     # func_sysd = env.transition_model.func()
     # func_obs = obsmodel.func()
     # jac_sysd = env.transition_model.jac_dx()
     # jac_sysd_u = env.transition_model.jac_du()
     # jac_obs = obsmodel.jac_dx()
-    
+
     # noise_sysd = pomdp_py.Gaussian([0,0], [[0.1,0],
     #                                        [0,0.1]]).random()
     # noise_obs_cov = obsmodel.noise_covariance(bt[0])
diff --git a/pomdp_problems/load_unload/README.rst b/pomdp_py/problems/load_unload/README.rst
similarity index 100%
rename from pomdp_problems/load_unload/README.rst
rename to pomdp_py/problems/load_unload/README.rst
diff --git a/pomdp_problems/load_unload/__init__.py b/pomdp_py/problems/load_unload/__init__.py
similarity index 100%
rename from pomdp_problems/load_unload/__init__.py
rename to pomdp_py/problems/load_unload/__init__.py
diff --git a/pomdp_problems/load_unload/load_unload.py b/pomdp_py/problems/load_unload/load_unload.py
similarity index 100%
rename from pomdp_problems/load_unload/load_unload.py
rename to pomdp_py/problems/load_unload/load_unload.py
diff --git a/pomdp_problems/maze/README.rst b/pomdp_py/problems/maze/README.rst
similarity index 100%
rename from pomdp_problems/maze/README.rst
rename to pomdp_py/problems/maze/README.rst
diff --git a/pomdp_problems/maze/__init__.py b/pomdp_py/problems/maze/__init__.py
similarity index 100%
rename from pomdp_problems/maze/__init__.py
rename to pomdp_py/problems/maze/__init__.py
diff --git a/pomdp_problems/maze/domain/__init__.py b/pomdp_py/problems/maze/domain/__init__.py
similarity index 100%
rename from pomdp_problems/maze/domain/__init__.py
rename to pomdp_py/problems/maze/domain/__init__.py
diff --git a/pomdp_problems/maze/domain/action.py b/pomdp_py/problems/maze/domain/action.py
similarity index 78%
rename from pomdp_problems/maze/domain/action.py
rename to pomdp_py/problems/maze/domain/action.py
index 2f98c228..b96a2ca4 100644
--- a/pomdp_problems/maze/domain/action.py
+++ b/pomdp_py/problems/maze/domain/action.py
@@ -4,5 +4,5 @@
 
 # Reuses the actions in the multi object search domain
 import pomdp_py
-from pomdp_problems.multi_object_search.domain.action\
+from pomdp_py.problems.multi_object_search.domain.action\
     import MotionAction, MoveForward, MoveBackward, MoveLeft, MoveRight, LookAction
diff --git a/pomdp_problems/maze/domain/observation.py b/pomdp_py/problems/maze/domain/observation.py
similarity index 100%
rename from pomdp_problems/maze/domain/observation.py
rename to pomdp_py/problems/maze/domain/observation.py
diff --git a/pomdp_problems/maze/domain/state.py b/pomdp_py/problems/maze/domain/state.py
similarity index 100%
rename from pomdp_problems/maze/domain/state.py
rename to pomdp_py/problems/maze/domain/state.py
diff --git a/pomdp_problems/maze/env/__init__.py b/pomdp_py/problems/maze/env/__init__.py
similarity index 100%
rename from pomdp_problems/maze/env/__init__.py
rename to pomdp_py/problems/maze/env/__init__.py
diff --git a/pomdp_problems/maze/env/env.py b/pomdp_py/problems/maze/env/env.py
similarity index 100%
rename from pomdp_problems/maze/env/env.py
rename to pomdp_py/problems/maze/env/env.py
diff --git a/pomdp_problems/maze/models/__init__.py b/pomdp_py/problems/maze/models/__init__.py
similarity index 100%
rename from pomdp_problems/maze/models/__init__.py
rename to pomdp_py/problems/maze/models/__init__.py
diff --git a/pomdp_problems/maze/models/components/__init__.py b/pomdp_py/problems/maze/models/components/__init__.py
similarity index 100%
rename from pomdp_problems/maze/models/components/__init__.py
rename to pomdp_py/problems/maze/models/components/__init__.py
diff --git a/pomdp_problems/maze/models/components/map.py b/pomdp_py/problems/maze/models/components/map.py
similarity index 100%
rename from pomdp_problems/maze/models/components/map.py
rename to pomdp_py/problems/maze/models/components/map.py
diff --git a/pomdp_problems/multi_object_search/README.rst b/pomdp_py/problems/multi_object_search/README.rst
similarity index 100%
rename from pomdp_problems/multi_object_search/README.rst
rename to pomdp_py/problems/multi_object_search/README.rst
diff --git a/pomdp_problems/multi_object_search/__init__.py b/pomdp_py/problems/multi_object_search/__init__.py
similarity index 100%
rename from pomdp_problems/multi_object_search/__init__.py
rename to pomdp_py/problems/multi_object_search/__init__.py
diff --git a/pomdp_problems/multi_object_search/agent/__init__.py b/pomdp_py/problems/multi_object_search/agent/__init__.py
similarity index 100%
rename from pomdp_problems/multi_object_search/agent/__init__.py
rename to pomdp_py/problems/multi_object_search/agent/__init__.py
diff --git a/pomdp_problems/multi_object_search/agent/agent.py b/pomdp_py/problems/multi_object_search/agent/agent.py
similarity index 100%
rename from pomdp_problems/multi_object_search/agent/agent.py
rename to pomdp_py/problems/multi_object_search/agent/agent.py
diff --git a/pomdp_problems/multi_object_search/agent/belief.py b/pomdp_py/problems/multi_object_search/agent/belief.py
similarity index 100%
rename from pomdp_problems/multi_object_search/agent/belief.py
rename to pomdp_py/problems/multi_object_search/agent/belief.py
diff --git a/pomdp_problems/multi_object_search/domain/__init__.py b/pomdp_py/problems/multi_object_search/domain/__init__.py
similarity index 100%
rename from pomdp_problems/multi_object_search/domain/__init__.py
rename to pomdp_py/problems/multi_object_search/domain/__init__.py
diff --git a/pomdp_problems/multi_object_search/domain/action.py b/pomdp_py/problems/multi_object_search/domain/action.py
similarity index 100%
rename from pomdp_problems/multi_object_search/domain/action.py
rename to pomdp_py/problems/multi_object_search/domain/action.py
diff --git a/pomdp_problems/multi_object_search/domain/observation.py b/pomdp_py/problems/multi_object_search/domain/observation.py
similarity index 100%
rename from pomdp_problems/multi_object_search/domain/observation.py
rename to pomdp_py/problems/multi_object_search/domain/observation.py
diff --git a/pomdp_problems/multi_object_search/domain/state.py b/pomdp_py/problems/multi_object_search/domain/state.py
similarity index 100%
rename from pomdp_problems/multi_object_search/domain/state.py
rename to pomdp_py/problems/multi_object_search/domain/state.py
diff --git a/pomdp_problems/multi_object_search/env/__init__.py b/pomdp_py/problems/multi_object_search/env/__init__.py
similarity index 100%
rename from pomdp_problems/multi_object_search/env/__init__.py
rename to pomdp_py/problems/multi_object_search/env/__init__.py
diff --git a/pomdp_problems/multi_object_search/env/env.py b/pomdp_py/problems/multi_object_search/env/env.py
similarity index 94%
rename from pomdp_problems/multi_object_search/env/env.py
rename to pomdp_py/problems/multi_object_search/env/env.py
index 8c375a8a..e04ab387 100644
--- a/pomdp_problems/multi_object_search/env/env.py
+++ b/pomdp_py/problems/multi_object_search/env/env.py
@@ -2,10 +2,10 @@
 
 import pomdp_py
 import copy
-from pomdp_problems.multi_object_search.models.transition_model import *
-from pomdp_problems.multi_object_search.models.reward_model import *
-from pomdp_problems.multi_object_search.models.components.sensor import *
-from pomdp_problems.multi_object_search.domain.state import *
+from pomdp_py.problems.multi_object_search.models.transition_model import *
+from pomdp_py.problems.multi_object_search.models.reward_model import *
+from pomdp_py.problems.multi_object_search.models.components.sensor import *
+from pomdp_py.problems.multi_object_search.domain.state import *
 
 class MosEnvironment(pomdp_py.Environment):
     """"""
@@ -33,7 +33,7 @@ def __init__(self, dim, init_state, sensors, obstacles=set({})):
         super().__init__(init_state,
                          transition_model,
                          reward_model)
-        
+
     @property
     def robot_ids(self):
         return set(self.sensors.keys())
@@ -61,14 +61,14 @@ def state_transition(self, action, execute=True, robot_id=None):
         next_state = copy.deepcopy(self.state)
         next_state.object_states[robot_id] =\
             self.transition_model[robot_id].sample(self.state, action)
-        
+
         reward = self.reward_model.sample(self.state, action, next_state,
                                           robot_id=robot_id)
         if execute:
             self.apply_transition(next_state)
             return reward
         else:
-            return next_state, reward        
+            return next_state, reward
 
 #### Interpret string as an initial world state ####
 def interpret(worldstr):
@@ -77,7 +77,7 @@ def interpret(worldstr):
     and returns the corresponding MosEnvironment.
 
     For example: This string
-    
+
     .. code-block:: text
 
         rx...
@@ -85,7 +85,7 @@ def interpret(worldstr):
         .....
         ***
         r: laser fov=90 min_range=1 max_range=10
-    
+
     describes a 3 by 5 world where x indicates obsticles and T indicates
     the "target object". T could be replaced by any upper-case letter A-Z
     which will serve as the object's id. Lower-case letters a-z (except for x)
@@ -93,14 +93,14 @@ def interpret(worldstr):
 
     After the world, the :code:`***` signals description of the sensor for each robot.
     For example "r laser 90 1 10" means that robot `r` will have a Laser2Dsensor
-    with fov 90, min_range 1.0, and max_range of 10.0.    
+    with fov 90, min_range 1.0, and max_range of 10.0.
 
     Args:
         worldstr (str): a string that describes the initial state of the world.
 
     Returns:
         MosEnvironment: the corresponding environment for the world description.
-            
+
     """
     worldlines = []
     sensorlines = []
@@ -115,11 +115,11 @@ def interpret(worldstr):
                 worldlines.append(line)
             if mode == "sensor":
                 sensorlines.append(line)
-    
+
     lines = [line for line in worldlines
              if len(line) > 0]
     w, l = len(worldlines[0]), len(worldlines)
-    
+
     objects = {}    # objid -> object_state(pose)
     obstacles = set({})  # objid
     robots = {}  # robot_id -> robot_state(pose)
@@ -137,12 +137,12 @@ def interpret(worldstr):
                 objid = 1000 + len(obstacles)  # obstacle id
                 objects[objid] = ObjectState(objid, "obstacle", (x,y))
                 obstacles.add(objid)
-                
+
             elif c.isupper():
                 # target object
                 objid = len(objects)
                 objects[objid] = ObjectState(objid, "target", (x,y))
-                
+
             elif c.islower():
                 # robot
                 robot_id = interpret_robot_id(c)
@@ -162,7 +162,7 @@ def interpret(worldstr):
         robot_name = line.split(":")[0].strip()
         robot_id = interpret_robot_id(robot_name)
         assert robot_id in robots, "Sensor specified for unknown robot %s" % (robot_name)
-        
+
         sensor_setting = line.split(":")[1].strip()
         sensor_type = sensor_setting.split(" ")[0].strip()
         sensor_params = {}
@@ -170,7 +170,7 @@ def interpret(worldstr):
             param_name = token.split("=")[0].strip()
             param_value = eval(token.split("=")[1].strip())
             sensor_params[param_name] = param_value
-        
+
         if sensor_type == "laser":
             sensor = Laser2DSensor(robot_id, **sensor_params)
         elif sensor_type == "proximity":
diff --git a/pomdp_problems/multi_object_search/env/visual.py b/pomdp_py/problems/multi_object_search/env/visual.py
similarity index 97%
rename from pomdp_problems/multi_object_search/env/visual.py
rename to pomdp_py/problems/multi_object_search/env/visual.py
index 41ccc5d2..40a3d6e3 100644
--- a/pomdp_problems/multi_object_search/env/visual.py
+++ b/pomdp_py/problems/multi_object_search/env/visual.py
@@ -12,11 +12,11 @@
 import numpy as np
 import random
 import pomdp_py.utils as util
-from pomdp_problems.multi_object_search.env.env import *
-from pomdp_problems.multi_object_search.domain.observation import *
-from pomdp_problems.multi_object_search.domain.action import *
-from pomdp_problems.multi_object_search.domain.state import *
-from pomdp_problems.multi_object_search.example_worlds import *
+from pomdp_py.problems.multi_object_search.env.env import *
+from pomdp_py.problems.multi_object_search.domain.observation import *
+from pomdp_py.problems.multi_object_search.domain.action import *
+from pomdp_py.problems.multi_object_search.domain.state import *
+from pomdp_py.problems.multi_object_search.example_worlds import *
 
 # Deterministic way to get object color
 def object_color(objid, count):
diff --git a/pomdp_problems/multi_object_search/example_worlds.py b/pomdp_py/problems/multi_object_search/example_worlds.py
similarity index 100%
rename from pomdp_problems/multi_object_search/example_worlds.py
rename to pomdp_py/problems/multi_object_search/example_worlds.py
diff --git a/pomdp_problems/multi_object_search/models/__init__.py b/pomdp_py/problems/multi_object_search/models/__init__.py
similarity index 100%
rename from pomdp_problems/multi_object_search/models/__init__.py
rename to pomdp_py/problems/multi_object_search/models/__init__.py
diff --git a/pomdp_problems/multi_object_search/models/components/__init__.py b/pomdp_py/problems/multi_object_search/models/components/__init__.py
similarity index 100%
rename from pomdp_problems/multi_object_search/models/components/__init__.py
rename to pomdp_py/problems/multi_object_search/models/components/__init__.py
diff --git a/pomdp_problems/multi_object_search/models/components/grid_map.py b/pomdp_py/problems/multi_object_search/models/components/grid_map.py
similarity index 88%
rename from pomdp_problems/multi_object_search/models/components/grid_map.py
rename to pomdp_py/problems/multi_object_search/models/components/grid_map.py
index 015e258a..38093b48 100644
--- a/pomdp_problems/multi_object_search/models/components/grid_map.py
+++ b/pomdp_py/problems/multi_object_search/models/components/grid_map.py
@@ -1,8 +1,8 @@
 """Optional grid map to assist collision avoidance during planning."""
 
-from pomdp_problems.multi_object_search.models.transition_model import RobotTransitionModel
-from pomdp_problems.multi_object_search.domain.action import *
-from pomdp_problems.multi_object_search.domain.state import *
+from pomdp_py.problems.multi_object_search.models.transition_model import RobotTransitionModel
+from pomdp_py.problems.multi_object_search.domain.action import *
+from pomdp_py.problems.multi_object_search.domain.state import *
 
 class GridMap:
     """This map assists the agent to avoid planning invalid
@@ -15,7 +15,7 @@ def __init__(self, width, length, obstacles):
             obstacles (dict): Map from objid to (x,y); The object is
                                    supposed to be an obstacle.
             width (int): width of the grid map
-            length (int): length of the grid map 
+            length (int): length of the grid map
         """
         self.width = width
         self.length = length
@@ -29,7 +29,7 @@ def __init__(self, width, length, obstacles):
         }
         # set of obstacle poses
         self.obstacle_poses = set({self._obstacles[objid]
-                                   for objid in self._obstacles})        
+                                   for objid in self._obstacles})
 
     def valid_motions(self, robot_id, robot_pose, all_motion_actions):
         """
diff --git a/pomdp_problems/multi_object_search/models/components/sensor.py b/pomdp_py/problems/multi_object_search/models/components/sensor.py
similarity index 98%
rename from pomdp_problems/multi_object_search/models/components/sensor.py
rename to pomdp_py/problems/multi_object_search/models/components/sensor.py
index 1a03a8b4..033956f2 100644
--- a/pomdp_problems/multi_object_search/models/components/sensor.py
+++ b/pomdp_py/problems/multi_object_search/models/components/sensor.py
@@ -2,8 +2,8 @@
 
 import math
 import numpy as np
-from pomdp_problems.multi_object_search.domain.action import *
-from pomdp_problems.multi_object_search.domain.observation import *
+from pomdp_py.problems.multi_object_search.domain.action import *
+from pomdp_py.problems.multi_object_search.domain.observation import *
 
 # Note that the occlusion of an object is implemented based on
 # whether a beam will hit an obstacle or some other object before
diff --git a/pomdp_problems/multi_object_search/models/observation_model.py b/pomdp_py/problems/multi_object_search/models/observation_model.py
similarity index 98%
rename from pomdp_problems/multi_object_search/models/observation_model.py
rename to pomdp_py/problems/multi_object_search/models/observation_model.py
index 9f5351fd..424118e9 100644
--- a/pomdp_problems/multi_object_search/models/observation_model.py
+++ b/pomdp_py/problems/multi_object_search/models/observation_model.py
@@ -21,9 +21,9 @@
 import math
 import random
 import numpy as np
-from pomdp_problems.multi_object_search.domain.state import *
-from pomdp_problems.multi_object_search.domain.action import *
-from pomdp_problems.multi_object_search.domain.observation import *
+from pomdp_py.problems.multi_object_search.domain.state import *
+from pomdp_py.problems.multi_object_search.domain.action import *
+from pomdp_py.problems.multi_object_search.domain.observation import *
 
 #### Observation Models ####
 class MosObservationModel(pomdp_py.OOObservationModel):
diff --git a/pomdp_problems/multi_object_search/models/policy_model.py b/pomdp_py/problems/multi_object_search/models/policy_model.py
similarity index 93%
rename from pomdp_problems/multi_object_search/models/policy_model.py
rename to pomdp_py/problems/multi_object_search/models/policy_model.py
index e21dd6a0..d0011576 100644
--- a/pomdp_problems/multi_object_search/models/policy_model.py
+++ b/pomdp_py/problems/multi_object_search/models/policy_model.py
@@ -1,11 +1,11 @@
-"""Policy model for 2D Multi-Object Search domain. 
+"""Policy model for 2D Multi-Object Search domain.
 It is optional for the agent to be equipped with an occupancy
 grid map of the environment.
 """
 
 import pomdp_py
 import random
-from pomdp_problems.multi_object_search.domain.action import *
+from pomdp_py.problems.multi_object_search.domain.action import *
 
 class PolicyModel(pomdp_py.RolloutPolicy):
     """Simple policy model. All actions are possible at any state."""
@@ -17,7 +17,7 @@ def __init__(self, robot_id, grid_map=None):
 
     def sample(self, state, **kwargs):
         return random.sample(self._get_all_actions(**kwargs), 1)[0]
-    
+
     def probability(self, action, state, **kwargs):
         raise NotImplementedError
 
diff --git a/pomdp_problems/multi_object_search/models/reward_model.py b/pomdp_py/problems/multi_object_search/models/reward_model.py
similarity index 96%
rename from pomdp_problems/multi_object_search/models/reward_model.py
rename to pomdp_py/problems/multi_object_search/models/reward_model.py
index 3791dbbb..4adc31bc 100644
--- a/pomdp_problems/multi_object_search/models/reward_model.py
+++ b/pomdp_py/problems/multi_object_search/models/reward_model.py
@@ -1,7 +1,7 @@
 """Reward model for 2D Multi-object Search domain"""
 
 import pomdp_py
-from pomdp_problems.multi_object_search.domain.action import *
+from pomdp_py.problems.multi_object_search.domain.action import *
 
 class MosRewardModel(pomdp_py.RewardModel):
     def __init__(self, target_objects, big=1000, small=1, robot_id=None):
@@ -14,18 +14,18 @@ def __init__(self, target_objects, big=1000, small=1, robot_id=None):
         self.big = big
         self.small = small
         self._target_objects = target_objects
-        
+
     def probability(self, reward, state, action, next_state, normalized=False, **kwargs):
         if reward == self._reward_func(state, action):
             return 1.0
         else:
             return 0.0
-        
+
     def sample(self, state, action, next_state,
                normalized=False, robot_id=None):
         # deterministic
         return self._reward_func(state, action, next_state, robot_id=robot_id)
-    
+
     def argmax(self, state, action, next_state, normalized=False, robot_id=None):
         """Returns the most likely reward"""
         return self._reward_func(state, action, next_state, robot_id=robot_id)
@@ -39,14 +39,14 @@ def _reward_func(self, state, action, next_state, robot_id=None):
             assert self._robot_id is not None,\
                 "Reward must be computed with respect to one robot."
             robot_id = self._robot_id
-            
+
         reward = 0
 
         # If the robot has detected all objects
         if len(state.object_states[robot_id]['objects_found'])\
            == len(self._target_objects):
             return 0  # no reward or penalty; the task is finished.
-        
+
         if isinstance(action, MotionAction):
             reward = reward - self.small - action.distance_cost
         elif isinstance(action, LookAction):
@@ -66,4 +66,3 @@ def _reward_func(self, state, action, next_state, robot_id=None):
                     # Has new detection. Award.
                     reward += self.big
         return reward
-    
diff --git a/pomdp_problems/multi_object_search/models/transition_model.py b/pomdp_py/problems/multi_object_search/models/transition_model.py
similarity index 95%
rename from pomdp_problems/multi_object_search/models/transition_model.py
rename to pomdp_py/problems/multi_object_search/models/transition_model.py
index 1563fb89..4b0d28de 100644
--- a/pomdp_problems/multi_object_search/models/transition_model.py
+++ b/pomdp_py/problems/multi_object_search/models/transition_model.py
@@ -10,9 +10,9 @@
 """
 import pomdp_py
 import copy
-from pomdp_problems.multi_object_search.domain.state import *
-from pomdp_problems.multi_object_search.domain.observation import *
-from pomdp_problems.multi_object_search.domain.action import *
+from pomdp_py.problems.multi_object_search.domain.state import *
+from pomdp_py.problems.multi_object_search.domain.observation import *
+from pomdp_py.problems.multi_object_search.domain.action import *
 
 ####### Transition Model #######
 class MosTransitionModel(pomdp_py.OOTransitionModel):
@@ -58,16 +58,16 @@ def probability(self, next_object_state, state, action):
             return self._epsilon
         else:
             return 1.0 - self._epsilon
-    
+
     def sample(self, state, action):
         """Returns next_object_state"""
         return self.argmax(state, action)
-    
+
     def argmax(self, state, action):
         """Returns the most likely next object_state"""
         return copy.deepcopy(state.object_states[self._objid])
 
-    
+
 class RobotTransitionModel(pomdp_py.TransitionModel):
     """We assume that the robot control is perfect and transitions are deterministic."""
     def __init__(self, sensor, dim, epsilon=1e-9):
@@ -126,8 +126,8 @@ def argmax(self, state, action):
             robot_state = state.object_states[self._robot_id]
 
         next_robot_state = copy.deepcopy(robot_state)
-        # camera direction is only not None when looking        
-        next_robot_state['camera_direction'] = None 
+        # camera direction is only not None when looking
+        next_robot_state['camera_direction'] = None
         if isinstance(action, MotionAction):
             # motion action
             next_robot_state['pose'] = \
@@ -151,7 +151,7 @@ def argmax(self, state, action):
             next_robot_state["objects_found"] = tuple(set(next_robot_state['objects_found'])\
                                                       | set(observed_target_objects))
         return next_robot_state
-    
+
     def sample(self, state, action):
         """Returns next_robot_state"""
         return self.argmax(state, action)
diff --git a/pomdp_problems/multi_object_search/problem.py b/pomdp_py/problems/multi_object_search/problem.py
similarity index 97%
rename from pomdp_problems/multi_object_search/problem.py
rename to pomdp_py/problems/multi_object_search/problem.py
index 2bced9f4..c62f6582 100644
--- a/pomdp_problems/multi_object_search/problem.py
+++ b/pomdp_py/problems/multi_object_search/problem.py
@@ -3,12 +3,12 @@
 to actually define the POMDP problem for multi-object search.
 Then, solve it using POUCT or POMCP."""
 import pomdp_py
-from pomdp_problems.multi_object_search.env.env import *
-from pomdp_problems.multi_object_search.env.visual import *
-from pomdp_problems.multi_object_search.agent.agent import *
-from pomdp_problems.multi_object_search.example_worlds import *
-from pomdp_problems.multi_object_search.domain.observation import *
-from pomdp_problems.multi_object_search.models.components.grid_map import *
+from pomdp_py.problems.multi_object_search.env.env import *
+from pomdp_py.problems.multi_object_search.env.visual import *
+from pomdp_py.problems.multi_object_search.agent.agent import *
+from pomdp_py.problems.multi_object_search.example_worlds import *
+from pomdp_py.problems.multi_object_search.domain.observation import *
+from pomdp_py.problems.multi_object_search.models.components.grid_map import *
 import argparse
 import time
 import random
diff --git a/pomdp_problems/rocksample/README.rst b/pomdp_py/problems/rocksample/README.rst
similarity index 100%
rename from pomdp_problems/rocksample/README.rst
rename to pomdp_py/problems/rocksample/README.rst
diff --git a/pomdp_problems/rocksample/__init__.py b/pomdp_py/problems/rocksample/__init__.py
similarity index 100%
rename from pomdp_problems/rocksample/__init__.py
rename to pomdp_py/problems/rocksample/__init__.py
diff --git a/pomdp_problems/rocksample/cythonize/__init__.py b/pomdp_py/problems/rocksample/cythonize/__init__.py
similarity index 100%
rename from pomdp_problems/rocksample/cythonize/__init__.py
rename to pomdp_py/problems/rocksample/cythonize/__init__.py
diff --git a/pomdp_problems/rocksample/cythonize/rocksample_problem.pyx b/pomdp_py/problems/rocksample/cythonize/rocksample_problem.pyx
similarity index 100%
rename from pomdp_problems/rocksample/cythonize/rocksample_problem.pyx
rename to pomdp_py/problems/rocksample/cythonize/rocksample_problem.pyx
diff --git a/pomdp_problems/rocksample/cythonize/run_rocksample.py b/pomdp_py/problems/rocksample/cythonize/run_rocksample.py
similarity index 100%
rename from pomdp_problems/rocksample/cythonize/run_rocksample.py
rename to pomdp_py/problems/rocksample/cythonize/run_rocksample.py
diff --git a/pomdp_problems/rocksample/rocksample_problem.py b/pomdp_py/problems/rocksample/rocksample_problem.py
similarity index 100%
rename from pomdp_problems/rocksample/rocksample_problem.py
rename to pomdp_py/problems/rocksample/rocksample_problem.py
diff --git a/pomdp_problems/tag/README.rst b/pomdp_py/problems/tag/README.rst
similarity index 100%
rename from pomdp_problems/tag/README.rst
rename to pomdp_py/problems/tag/README.rst
diff --git a/pomdp_problems/tag/__init__.py b/pomdp_py/problems/tag/__init__.py
similarity index 100%
rename from pomdp_problems/tag/__init__.py
rename to pomdp_py/problems/tag/__init__.py
diff --git a/pomdp_problems/tag/agent/__init__.py b/pomdp_py/problems/tag/agent/__init__.py
similarity index 100%
rename from pomdp_problems/tag/agent/__init__.py
rename to pomdp_py/problems/tag/agent/__init__.py
diff --git a/pomdp_problems/tag/agent/agent.py b/pomdp_py/problems/tag/agent/agent.py
similarity index 84%
rename from pomdp_problems/tag/agent/agent.py
rename to pomdp_py/problems/tag/agent/agent.py
index cb6efa8d..dd84a67b 100644
--- a/pomdp_problems/tag/agent/agent.py
+++ b/pomdp_py/problems/tag/agent/agent.py
@@ -1,20 +1,20 @@
 import pomdp_py
 import copy
 import random
-from pomdp_problems.tag.domain.observation import *
-from pomdp_problems.tag.domain.action import *
-from pomdp_problems.tag.domain.state import *
-from pomdp_problems.tag.models.observation_model import *
-from pomdp_problems.tag.models.transition_model import *
-from pomdp_problems.tag.models.reward_model import *
-from pomdp_problems.tag.models.policy_model import *
-from pomdp_problems.tag.models.components.motion_policy import *
-from pomdp_problems.tag.models.components.grid_map import *
+from pomdp_py.problems.tag.domain.observation import *
+from pomdp_py.problems.tag.domain.action import *
+from pomdp_py.problems.tag.domain.state import *
+from pomdp_py.problems.tag.models.observation_model import *
+from pomdp_py.problems.tag.models.transition_model import *
+from pomdp_py.problems.tag.models.reward_model import *
+from pomdp_py.problems.tag.models.policy_model import *
+from pomdp_py.problems.tag.models.components.motion_policy import *
+from pomdp_py.problems.tag.models.components.grid_map import *
 
 ## initialize belief
 def initialize_belief(grid_map, init_robot_position, prior={}):
     """Initialize belief.
-    
+
     Args:
         grid_map (GridMap): Holds information of the map occupancy
         prior (dict): A map from (x,y)->[0,1]. If empty, the belief
@@ -47,7 +47,7 @@ def initialize_belief(grid_map, init_robot_position, prior={}):
 
 def initialize_particles_belief(grid_map, init_robot_position, num_particles=100, prior={}):
     """Initialize belief.
-    
+
     Args:
         grid_map (GridMap): Holds information of the map occupancy
         prior (dict): A map from (x,y)->[0,1]. If empty, the belief
@@ -60,14 +60,14 @@ def initialize_particles_belief(grid_map, init_robot_position, num_particles=100
             state = TagState(init_robot_position, pose)
             amount_to_add = (prior[objid][pose] / prior_sum) * num_particles
             for _ in range(amount_to_add):
-                particles.append(state)    
+                particles.append(state)
     else:
         while len(particles) < num_particles:
             target_position = (random.randint(0, grid_map.width-1),
                                random.randint(0, grid_map.length-1))
             if target_position in grid_map.obstacle_poses:
                 # Skip obstacles
-                continue            
+                continue
             state = TagState(init_robot_position, target_position, False)
             particles.append(state)
     return pomdp_py.Particles(particles)
@@ -78,7 +78,7 @@ def belief_update(agent, real_action, real_observation):
     # Update agent belief
     current_mpe_state = agent.cur_belief.mpe()
     next_robot_position = agent.transition_model.sample(current_mpe_state, real_action).robot_position
-    
+
     next_state_space = set({})
     for state in agent.cur_belief:
         next_state = copy.deepcopy(state)
@@ -89,20 +89,20 @@ def belief_update(agent, real_action, real_observation):
         agent.cur_belief, real_action, real_observation,
         agent.observation_model, agent.transition_model,
         next_state_space=next_state_space)
-    
+
     agent.set_belief(new_belief)
 
 class TagAgent(pomdp_py.Agent):
-    
+
     def __init__(self,
                  init_belief,
                  grid_map,
                  pr_stay=0.2,
                  small=1,
                  big=10):
-        self._grid_map = grid_map        
+        self._grid_map = grid_map
         target_motion_policy = TagTargetMotionPolicy(grid_map,
-                                                     pr_stay)        
+                                                     pr_stay)
         transition_model = TagTransitionModel(grid_map,
                                               target_motion_policy)
         reward_model = TagRewardModel(small=small, big=big)
@@ -117,5 +117,3 @@ def __init__(self,
     def clear_history(self):
         """Custum function; clear history"""
         self._history = None
-        
-        
diff --git a/pomdp_problems/tag/constants.py b/pomdp_py/problems/tag/constants.py
similarity index 100%
rename from pomdp_problems/tag/constants.py
rename to pomdp_py/problems/tag/constants.py
diff --git a/pomdp_problems/tag/domain/__init__.py b/pomdp_py/problems/tag/domain/__init__.py
similarity index 100%
rename from pomdp_problems/tag/domain/__init__.py
rename to pomdp_py/problems/tag/domain/__init__.py
diff --git a/pomdp_problems/tag/domain/action.py b/pomdp_py/problems/tag/domain/action.py
similarity index 84%
rename from pomdp_problems/tag/domain/action.py
rename to pomdp_py/problems/tag/domain/action.py
index 72822e9d..04694473 100644
--- a/pomdp_problems/tag/domain/action.py
+++ b/pomdp_py/problems/tag/domain/action.py
@@ -6,11 +6,11 @@
 
 # Reuses the actions in the multi object search domain
 import pomdp_py
-from pomdp_problems.multi_object_search.domain.action\
+from pomdp_py.problems.multi_object_search.domain.action\
     import Action, MotionAction, MoveEast2D, MoveWest2D, MoveSouth2D, MoveNorth2D
 
 MOTION_ACTIONS = {MoveEast2D, MoveWest2D, MoveSouth2D, MoveNorth2D}
 
 class TagAction(Action):
     def __init__(self):
-        super().__init__("tag")    
+        super().__init__("tag")
diff --git a/pomdp_problems/tag/domain/observation.py b/pomdp_py/problems/tag/domain/observation.py
similarity index 100%
rename from pomdp_problems/tag/domain/observation.py
rename to pomdp_py/problems/tag/domain/observation.py
diff --git a/pomdp_problems/tag/domain/state.py b/pomdp_py/problems/tag/domain/state.py
similarity index 100%
rename from pomdp_problems/tag/domain/state.py
rename to pomdp_py/problems/tag/domain/state.py
diff --git a/pomdp_problems/tag/env/__init__.py b/pomdp_py/problems/tag/env/__init__.py
similarity index 100%
rename from pomdp_problems/tag/env/__init__.py
rename to pomdp_py/problems/tag/env/__init__.py
diff --git a/pomdp_problems/tag/env/env.py b/pomdp_py/problems/tag/env/env.py
similarity index 78%
rename from pomdp_problems/tag/env/env.py
rename to pomdp_py/problems/tag/env/env.py
index 3fafe047..a66ac829 100644
--- a/pomdp_problems/tag/env/env.py
+++ b/pomdp_py/problems/tag/env/env.py
@@ -1,11 +1,11 @@
 import pomdp_py
-from pomdp_problems.tag.domain.state import *
-from pomdp_problems.tag.models.transition_model import *
-from pomdp_problems.tag.models.reward_model import *
-from pomdp_problems.tag.models.components.motion_policy import *
-from pomdp_problems.tag.models.components.grid_map import *
-from pomdp_problems.multi_object_search.env.env import interpret
-from pomdp_problems.multi_object_search.env.visual import MosViz
+from pomdp_py.problems.tag.domain.state import *
+from pomdp_py.problems.tag.models.transition_model import *
+from pomdp_py.problems.tag.models.reward_model import *
+from pomdp_py.problems.tag.models.components.motion_policy import *
+from pomdp_py.problems.tag.models.components.grid_map import *
+from pomdp_py.problems.multi_object_search.env.env import interpret
+from pomdp_py.problems.multi_object_search.env.visual import MosViz
 
 class TagEnvironment(pomdp_py.Environment):
 
@@ -27,7 +27,7 @@ def __init__(self,
     @property
     def width(self):
         return self._grid_map.width
-    
+
     @property
     def length(self):
         return self._grid_map.length
@@ -48,7 +48,7 @@ def from_str(cls, worldstr, **kwargs):
                 targets.append(objid)
             else:
                 obstacle_poses.add(objects[objid].pose)
-        assert len(targets) == 1, "Does not support multiple objects."                        
+        assert len(targets) == 1, "Does not support multiple objects."
         target_position = objects[targets[0]].pose
         init_state = TagState(robot_position, target_position, False)
         grid_map = GridMap(dim[0], dim[1], obstacle_poses)
diff --git a/pomdp_problems/tag/env/visual.py b/pomdp_py/problems/tag/env/visual.py
similarity index 96%
rename from pomdp_problems/tag/env/visual.py
rename to pomdp_py/problems/tag/env/visual.py
index 67cc98b5..b8ee6d05 100644
--- a/pomdp_problems/tag/env/visual.py
+++ b/pomdp_py/problems/tag/env/visual.py
@@ -5,12 +5,12 @@
 import numpy as np
 import random
 import pomdp_py.utils as util
-from pomdp_problems.tag.env.env import *
-from pomdp_problems.tag.domain.observation import *
-from pomdp_problems.tag.domain.action import *
-from pomdp_problems.tag.domain.state import *
-from pomdp_problems.tag.example_worlds import *
-from pomdp_problems.tag.models.observation_model import *
+from pomdp_py.problems.tag.env.env import *
+from pomdp_py.problems.tag.domain.observation import *
+from pomdp_py.problems.tag.domain.action import *
+from pomdp_py.problems.tag.domain.state import *
+from pomdp_py.problems.tag.example_worlds import *
+from pomdp_py.problems.tag.models.observation_model import *
 
 #### Visualization through pygame ####
 class TagViz:
diff --git a/pomdp_problems/tag/example_worlds.py b/pomdp_py/problems/tag/example_worlds.py
similarity index 100%
rename from pomdp_problems/tag/example_worlds.py
rename to pomdp_py/problems/tag/example_worlds.py
diff --git a/pomdp_problems/tag/experiment.py b/pomdp_py/problems/tag/experiment.py
similarity index 96%
rename from pomdp_problems/tag/experiment.py
rename to pomdp_py/problems/tag/experiment.py
index be61938b..647ab5fd 100644
--- a/pomdp_problems/tag/experiment.py
+++ b/pomdp_py/problems/tag/experiment.py
@@ -1,6 +1,6 @@
 """Simple experiment to get mean"""
 
-from pomdp_problems.tag.problem import *
+from pomdp_py.problems.tag.problem import *
 import numpy as np
 
 def trial(worldstr, **kwargs):
@@ -8,7 +8,7 @@ def trial(worldstr, **kwargs):
     free_cells = grid_map.free_cells()
     init_robot_position = random.sample(free_cells, 1)[0]
     init_target_position = random.sample(free_cells, 1)[0]
-    
+
     problem = TagProblem(init_robot_position,
                          init_target_position,
                          grid_map, **kwargs)
@@ -40,5 +40,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-
-
diff --git a/pomdp_problems/tag/models/__init__.py b/pomdp_py/problems/tag/models/__init__.py
similarity index 100%
rename from pomdp_problems/tag/models/__init__.py
rename to pomdp_py/problems/tag/models/__init__.py
diff --git a/pomdp_problems/tag/models/components/__init__.py b/pomdp_py/problems/tag/models/components/__init__.py
similarity index 100%
rename from pomdp_problems/tag/models/components/__init__.py
rename to pomdp_py/problems/tag/models/components/__init__.py
diff --git a/pomdp_problems/tag/models/components/grid_map.py b/pomdp_py/problems/tag/models/components/grid_map.py
similarity index 87%
rename from pomdp_problems/tag/models/components/grid_map.py
rename to pomdp_py/problems/tag/models/components/grid_map.py
index 663e2276..462c481a 100644
--- a/pomdp_problems/tag/models/components/grid_map.py
+++ b/pomdp_py/problems/tag/models/components/grid_map.py
@@ -1,6 +1,6 @@
-from pomdp_problems.tag.domain.action import *
-from pomdp_problems.tag.models.transition_model import TagTransitionModel
-from pomdp_problems.multi_object_search.env.env import interpret
+from pomdp_py.problems.tag.domain.action import *
+from pomdp_py.problems.tag.models.transition_model import TagTransitionModel
+from pomdp_py.problems.multi_object_search.env.env import interpret
 
 class GridMap:
 
@@ -43,4 +43,3 @@ def free_cells(self):
                      for y in range(self.length)
                      if (x,y) not in self.obstacle_poses})
         return cells
-                
diff --git a/pomdp_problems/tag/models/components/motion_policy.py b/pomdp_py/problems/tag/models/components/motion_policy.py
similarity index 97%
rename from pomdp_problems/tag/models/components/motion_policy.py
rename to pomdp_py/problems/tag/models/components/motion_policy.py
index e20445cd..8900e423 100644
--- a/pomdp_problems/tag/models/components/motion_policy.py
+++ b/pomdp_py/problems/tag/models/components/motion_policy.py
@@ -1,8 +1,8 @@
 import pomdp_py
 import random
 from pomdp_py.utils.math import euclidean_dist
-import pomdp_problems.tag.constants as constants
-from pomdp_problems.tag.models.transition_model import TagTransitionModel
+import pomdp_py.problems.tag.constants as constants
+from pomdp_py.problems.tag.models.transition_model import TagTransitionModel
 
 
 class TagTargetMotionPolicy(pomdp_py.GenerativeDistribution):
diff --git a/pomdp_problems/tag/models/observation_model.py b/pomdp_py/problems/tag/models/observation_model.py
similarity index 92%
rename from pomdp_problems/tag/models/observation_model.py
rename to pomdp_py/problems/tag/models/observation_model.py
index 5ee836e2..eda84bb6 100644
--- a/pomdp_problems/tag/models/observation_model.py
+++ b/pomdp_py/problems/tag/models/observation_model.py
@@ -1,6 +1,6 @@
 import pomdp_py
-from pomdp_problems.tag.domain.observation import *
-import pomdp_problems.tag.constants as constants
+from pomdp_py.problems.tag.domain.observation import *
+import pomdp_py.problems.tag.constants as constants
 
 class TagObservationModel(pomdp_py.ObservationModel):
     """In this observation model, the robot deterministically
@@ -21,7 +21,7 @@ def probability(self, observation, next_state, action, **kwargs):
                 return 1.0 - constants.EPSILON
             else:
                 return constants.EPSILON
-    
+
     def sample(self, next_state, action):
         """There is no stochaisticity in the observation model"""
         if next_state.robot_position == next_state.target_position:
diff --git a/pomdp_problems/tag/models/policy_model.py b/pomdp_py/problems/tag/models/policy_model.py
similarity index 88%
rename from pomdp_problems/tag/models/policy_model.py
rename to pomdp_py/problems/tag/models/policy_model.py
index 7084ed1b..fcc89a61 100644
--- a/pomdp_problems/tag/models/policy_model.py
+++ b/pomdp_py/problems/tag/models/policy_model.py
@@ -1,7 +1,7 @@
 import pomdp_py
 import random
-from pomdp_problems.tag.domain.action import *
-from pomdp_problems.tag.models.transition_model import *
+from pomdp_py.problems.tag.domain.action import *
+from pomdp_py.problems.tag.models.transition_model import *
 
 class TagPolicyModel(pomdp_py.RolloutPolicy):
     def __init__(self, grid_map=None):
@@ -9,7 +9,7 @@ def __init__(self, grid_map=None):
 
     def sample(self, state, **kwargs):
         return random.sample(self.get_all_actions(state, **kwargs), 1)[0]
-    
+
     def get_all_actions(self, state=None, history=None):
         if state is not None:
             if self._grid_map is not None:
@@ -20,4 +20,3 @@ def get_all_actions(self, state=None, history=None):
 
     def rollout(self, state, history=None):
         return random.sample(self.get_all_actions(state=state, history=history), 1)[0]
-
diff --git a/pomdp_problems/tag/models/reward_model.py b/pomdp_py/problems/tag/models/reward_model.py
similarity index 94%
rename from pomdp_problems/tag/models/reward_model.py
rename to pomdp_py/problems/tag/models/reward_model.py
index 6cac6699..36499fd9 100644
--- a/pomdp_problems/tag/models/reward_model.py
+++ b/pomdp_py/problems/tag/models/reward_model.py
@@ -1,5 +1,5 @@
 import pomdp_py
-from pomdp_problems.tag.domain.action import *
+from pomdp_py.problems.tag.domain.action import *
 
 class TagRewardModel(pomdp_py.RewardModel):
 
@@ -12,7 +12,7 @@ def probability(self, reward, state, action, next_state, normalized=False, **kwa
             return 1.0
         else:
             return 0.0
-        
+
     def sample(self, state, action, next_state):
         # deterministic
         return self._reward_func(state, action, next_state)
diff --git a/pomdp_problems/tag/models/transition_model.py b/pomdp_py/problems/tag/models/transition_model.py
similarity index 97%
rename from pomdp_problems/tag/models/transition_model.py
rename to pomdp_py/problems/tag/models/transition_model.py
index 16e6e513..4a524f75 100644
--- a/pomdp_problems/tag/models/transition_model.py
+++ b/pomdp_py/problems/tag/models/transition_model.py
@@ -8,8 +8,8 @@
 """
 import copy
 import pomdp_py
-import pomdp_problems.tag.constants as constants
-from pomdp_problems.tag.domain.action import *
+import pomdp_py.problems.tag.constants as constants
+from pomdp_py.problems.tag.domain.action import *
 
 class TagTransitionModel(pomdp_py.TransitionModel):
 
diff --git a/pomdp_problems/tag/problem.py b/pomdp_py/problems/tag/problem.py
similarity index 96%
rename from pomdp_problems/tag/problem.py
rename to pomdp_py/problems/tag/problem.py
index 103b14a8..d5a9df14 100644
--- a/pomdp_problems/tag/problem.py
+++ b/pomdp_py/problems/tag/problem.py
@@ -1,9 +1,9 @@
 import pomdp_py
 import random
-from pomdp_problems.tag.env.env import *
-from pomdp_problems.tag.env.visual import *
-from pomdp_problems.tag.agent.agent import *
-from pomdp_problems.tag.example_worlds import *
+from pomdp_py.problems.tag.env.env import *
+from pomdp_py.problems.tag.env.visual import *
+from pomdp_py.problems.tag.agent.agent import *
+from pomdp_py.problems.tag.example_worlds import *
 import time
 
 class TagProblem(pomdp_py.POMDP):
diff --git a/pomdp_problems/tiger/README.rst b/pomdp_py/problems/tiger/README.rst
similarity index 100%
rename from pomdp_problems/tiger/README.rst
rename to pomdp_py/problems/tiger/README.rst
diff --git a/pomdp_py/problems/tiger/__init__.py b/pomdp_py/problems/tiger/__init__.py
new file mode 100644
index 00000000..f6863765
--- /dev/null
+++ b/pomdp_py/problems/tiger/__init__.py
@@ -0,0 +1 @@
+from .tiger_problem import TigerProblem, test_planner, make_tiger
diff --git a/pomdp_problems/tiger/cythonize/__init__.py b/pomdp_py/problems/tiger/cythonize/__init__.py
similarity index 100%
rename from pomdp_problems/tiger/cythonize/__init__.py
rename to pomdp_py/problems/tiger/cythonize/__init__.py
diff --git a/pomdp_problems/tiger/cythonize/run_tiger.py b/pomdp_py/problems/tiger/cythonize/run_tiger.py
similarity index 100%
rename from pomdp_problems/tiger/cythonize/run_tiger.py
rename to pomdp_py/problems/tiger/cythonize/run_tiger.py
diff --git a/pomdp_problems/tiger/cythonize/tiger_problem.pyx b/pomdp_py/problems/tiger/cythonize/tiger_problem.pyx
similarity index 100%
rename from pomdp_problems/tiger/cythonize/tiger_problem.pyx
rename to pomdp_py/problems/tiger/cythonize/tiger_problem.pyx
diff --git a/pomdp_problems/tiger/tiger_problem.py b/pomdp_py/problems/tiger/tiger_problem.py
similarity index 90%
rename from pomdp_problems/tiger/tiger_problem.py
rename to pomdp_py/problems/tiger/tiger_problem.py
index 141898e1..039f8846 100644
--- a/pomdp_problems/tiger/tiger_problem.py
+++ b/pomdp_py/problems/tiger/tiger_problem.py
@@ -39,6 +39,7 @@
 import random
 import numpy as np
 import sys
+import copy
 
 class TigerState(pomdp_py.State):
     def __init__(self, name):
@@ -245,9 +246,9 @@ def test_planner(tiger_problem, planner, nsteps=3,
             import pdb; pdb.set_trace()
 
         print("==== Step %d ====" % (i+1))
-        print("True state:", tiger_problem.env.state)
-        print("Belief:", tiger_problem.agent.cur_belief)
-        print("Action:", action)
+        print(f"True state: {tiger_problem.env.state}")
+        print(f"Belief: {tiger_problem.agent.cur_belief}")
+        print(f"Action: {action}")
         # There is no state transition for the tiger domain.
         # In general, the ennvironment state can be transitioned
         # using
@@ -297,41 +298,47 @@ def test_planner(tiger_problem, planner, nsteps=3,
             # until every time door is opened.
             print("\n")
 
+def make_tiger(noise=0.15, init_state="tiger-left", init_belief=[0.5, 0.5]):
+    """Convenient function to quickly build a tiger domain.
+    Useful for testing"""
+    tiger = TigerProblem(noise, TigerState(init_state),
+                         pomdp_py.Histogram({TigerState("tiger-left"): init_belief[0],
+                                             TigerState("tiger-right"): init_belief[1]}))
+    return tiger
+
+
 def main():
-    init_true_state = random.choice([TigerState("tiger-left"),
-                                     TigerState("tiger-right")])
+    init_true_state = random.choice(["tiger-left",
+                                     "tiger-right"])
     init_belief = pomdp_py.Histogram({TigerState("tiger-left"): 0.5,
                                       TigerState("tiger-right"): 0.5})
-    tiger_problem = TigerProblem(0.15,  # observation noise
-                                 init_true_state, init_belief)
+    tiger = make_tiger(init_state=init_true_state)
+    init_belief = tiger.agent.belief
 
     print("** Testing value iteration **")
     vi = pomdp_py.ValueIteration(horizon=3, discount_factor=0.95)
-    test_planner(tiger_problem, vi, nsteps=3)
-
-    # Reset agent belief
-    tiger_problem.agent.set_belief(init_belief, prior=True)
+    test_planner(tiger, vi, nsteps=3)
 
     print("\n** Testing POUCT **")
     pouct = pomdp_py.POUCT(max_depth=3, discount_factor=0.95,
                            num_sims=4096, exploration_const=50,
-                           rollout_policy=tiger_problem.agent.policy_model,
+                           rollout_policy=tiger.agent.policy_model,
                            show_progress=True)
-    test_planner(tiger_problem, pouct, nsteps=10)
-    TreeDebugger(tiger_problem.agent.tree).pp
+    test_planner(tiger, pouct, nsteps=10)
+    TreeDebugger(tiger.agent.tree).pp
 
     # Reset agent belief
-    tiger_problem.agent.set_belief(init_belief, prior=True)
-    tiger_problem.agent.tree = None
+    tiger.agent.set_belief(init_belief, prior=True)
+    tiger.agent.tree = None
 
     print("** Testing POMCP **")
-    tiger_problem.agent.set_belief(pomdp_py.Particles.from_histogram(init_belief, num_particles=100), prior=True)
+    tiger.agent.set_belief(pomdp_py.Particles.from_histogram(init_belief, num_particles=100), prior=True)
     pomcp = pomdp_py.POMCP(max_depth=3, discount_factor=0.95,
                            num_sims=1000, exploration_const=50,
-                           rollout_policy=tiger_problem.agent.policy_model,
+                           rollout_policy=tiger.agent.policy_model,
                            show_progress=True, pbar_update_interval=500)
-    test_planner(tiger_problem, pomcp, nsteps=10)
-    TreeDebugger(tiger_problem.agent.tree).pp
+    test_planner(tiger, pomcp, nsteps=10)
+    TreeDebugger(tiger.agent.tree).pp
 
 if __name__ == '__main__':
     main()
diff --git a/pomdp_py/utils/__init__.py b/pomdp_py/utils/__init__.py
index dc495f37..59d457f8 100644
--- a/pomdp_py/utils/__init__.py
+++ b/pomdp_py/utils/__init__.py
@@ -3,7 +3,6 @@
                                  R_x, R_y, R_z, R_between, T,
                                  to_radians, approx_equal)
 from pomdp_py.utils.misc import (remap, json_safe, safe_slice, similar, special_char)
-from pomdp_py.utils.plotting import (plot_points, save_plot)
 from pomdp_py.utils.colors import (lighter, rgb_to_hex, hex_to_rgb,
                                    inverse_color_rgb, inverse_color_hex,
                                    random_unique_color)
diff --git a/pomdp_py/utils/test_utils.py b/pomdp_py/utils/test_utils.py
deleted file mode 100644
index 78ac40e3..00000000
--- a/pomdp_py/utils/test_utils.py
+++ /dev/null
@@ -1,18 +0,0 @@
-import pomdp_py
-from pomdp_problems.tiger.tiger_problem import TigerProblem
-from pomdp_problems.tiger.tiger_problem import TigerState
-import os
-import glob
-
-def remove_files(pattern):
-    file_list = glob.glob(pattern)
-    for file_path in file_list:
-        os.remove(file_path)
-
-def make_tiger(noise=0.15, init_state="tiger-left", init_belief=[0.5, 0.5]):
-    """Convenient function to quickly build a tiger domain.
-    Useful for testing"""
-    tiger = TigerProblem(noise, TigerState(init_state),
-                         pomdp_py.Histogram({TigerState("tiger-left"): init_belief[0],
-                                             TigerState("tiger-right"): init_belief[1]}))
-    return tiger
diff --git a/pomdp_py/visual/__init__.py b/pomdp_py/visual/__init__.py
deleted file mode 100644
index 8b137891..00000000
--- a/pomdp_py/visual/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-
diff --git a/pomdp_py/visual/search_tree.py b/pomdp_py/visual/search_tree.py
deleted file mode 100644
index 34ccf37f..00000000
--- a/pomdp_py/visual/search_tree.py
+++ /dev/null
@@ -1,143 +0,0 @@
-################################################################################
-# WARNING: The library does not require the installation of networkx and
-# pygraphviz any more. But this code will be kept. If you would like to plot the
-# MCTS search tree, you should install networkx and pygraphviz yourself.
-################################################################################
-
-import networkx as nx
-from networkx.drawing.nx_agraph import graphviz_layout
-import matplotlib.pyplot as plt
-from pomdp_py import QNode, VNode, RootVNode, util
-
-# ---- POUCT Visualization ---- #
-
-def _build_graph(G, root, parent, conn, depth,
-                 max_depth=None, visit_threshold=0,
-                 max_value=None, min_value=None,
-                 relabel_actions={},
-                 relabel_observations={}):
-    if max_depth is not None and depth >= max_depth:
-        return
-    G.add_node(root)
-    if parent is not None:
-        if isinstance(root, VNode):
-            if len(relabel_observations) > 0:
-                conn = relabel_observations[conn]
-        else:
-            if len(relabel_actions) > 0:
-                conn = relabel_actions[conn]
-        if max_value - min_value > 0.0:
-            weight = util.remap(root.value,
-                                min_value, max_value,
-                                0.1, 3.0)
-        else:
-            weight = 0.1
-        G.add_edge(parent, root, **{"label":conn,
-                                    "weight": weight})
-    min_value = float('inf')
-    max_value = float('-inf')
-    for c in root.children:
-        max_value = max(max_value, root[c].value)
-        min_value = min(min_value, root[c].value)
-    for c in root.children:
-        if root[c].num_visits > visit_threshold:
-            _build_graph(G, root[c], root, c, depth+1,
-                         max_value=max_value, min_value=min_value,
-                         max_depth=max_depth, visit_threshold=visit_threshold,
-                         relabel_actions=relabel_actions, relabel_observations=relabel_observations)
-
-def _build_relabel_dict(root, conn, depth, actions, observations, max_depth=None, visit_threshold=0):
-    """Traverse the tree and collect unique actions and observations,
-    store them in dictionaries `actions` and `observations`. Requires
-    all observations and actions to be hashable."""
-    if max_depth is not None and depth >= max_depth:
-        return
-    if conn is not None:
-        if isinstance(root, VNode):
-            if conn not in observations:
-                observations[conn] = "o" + str(len(observations))
-        else:
-            if conn not in actions:
-                actions[conn] = "a" + str(len(actions))
-    for c in root.children:
-        if root[c].num_visits > visit_threshold:
-            _build_relabel_dict(root[c], c,  depth+1, actions, observations, max_depth=max_depth, visit_threshold=visit_threshold)
-
-def visualize_pouct_search_tree(root, max_depth=1,
-                                visit_threshold=1, anonymize=False,
-                                anonymize_actions=False, anonymize_observations=False,
-                                output_file=None, use_dot=False, ax=None):
-    """
-    Visualize the given tree up to depth `max_depth`. Display nodes
-    with number of visits >= `visit_threshold`.
-
-    Caveat: This only works well if the search tree depth is shallow.
-    For larger trees, please use a combination of tree debugger (utils.debugging.TreeDebugger)
-    and visualizer (custom in heritance of visual.visualizer.Visualizer).
-
-    If anonymize is True, will only display actions as a1,a2,... and observations
-    as o1,o2,... .
-    """
-    relabel_actions = {}
-    relabel_observations = {}
-    if anonymize:
-        anonymize_actions = True
-        anonymize_observations = True
-    if anonymize_actions or anonymize_observations:
-        _build_relabel_dict(root, None, 0, relabel_actions, relabel_observations,
-                            max_depth=max_depth, visit_threshold=visit_threshold)
-
-        print("---- Action labels ----")
-        action_map = {relabel_actions[action]:action for action in relabel_actions}
-        for label in sorted(action_map):
-            print("%s : %s" % (label, action_map[label]))
-        print("---- Observation labels ----")
-        observation_map = {relabel_observations[ob]:ob for ob in relabel_observations}
-        for label in sorted(observation_map):
-            print("%s : %s" % (label, observation_map[label]))
-
-    if not anonymize_actions:
-        relabel_actions = {}
-    if not anonymize_observations:
-        relabel_observations = {}
-
-    # Build a networkx graph.
-    G = nx.DiGraph()
-    _build_graph(G, root, None, None, 0, max_depth=max_depth, visit_threshold=visit_threshold,
-                 relabel_actions=relabel_actions, relabel_observations=relabel_observations)
-    if use_dot:
-        if output_file is None:
-            raise TypeError("Please provide output .dot file.")
-        nx.nx_agraph.write_dot(G, output_file)
-        print("Dot file saved at %s" % output_file)
-        print("Please run `dot -Tpng %s > %s.png" % (output_file, output_file))
-    else:
-        node_labels = {}
-        color_map = []
-        for node in G.nodes():
-            belief_str = ""
-            if hasattr(node, "belief"):
-                belief_str = " | %d" % len(node.belief)
-            if isinstance(node, RootVNode):
-                color_map.append("cyan")
-                node_labels[node] = "R(%d | %.2f%s)" % (node.num_visits, node.value, belief_str)
-            elif isinstance(node, VNode):
-                color_map.append("yellow")
-                node_labels[node] = "V(%d | %.2f%s)" % (node.num_visits, node.value, belief_str)
-            else:
-                color_map.append("orange")
-                node_labels[node] = "Q(%d | %.2f)" % (node.num_visits, node.value)
-        edge_labels = {(edge[0],edge[1]): edge[2]["label"] for edge in G.edges(data=True)}
-        edge_widths = [edge[2]["weight"] for edge in G.edges(data=True)]
-
-        pos = graphviz_layout(G, prog='dot')
-        nx.draw_networkx(G, pos,
-                         node_color=color_map, labels=node_labels,
-                         width=edge_widths,
-                         font_size=7, ax=ax)
-        nx.draw_networkx_edge_labels(G, pos,
-                                     edge_labels=edge_labels, ax=ax)
-        if output_file is None:
-            plt.show()
-        else:
-            plt.savefig(output_file)
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..317b1555
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,39 @@
+[build-system]
+requires = ["setuptools", "wheel", "Cython>=3.0.8"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "pomdp-py"
+version = "1.3.4"
+description = "Python POMDP Library."
+authors = [{name = "Kaiyu Zheng", email = "kzheng10@cs.brown.edu"}]
+license = {file = "LICENSE"}
+readme = "README.rst"
+keywords = ["Partially Observable Markov Decision Process", "POMDP"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python",
+    "Programming Language :: Cython",
+    "Topic :: Scientific/Engineering"
+
+]
+dependencies = [
+    "numpy>=1.18.0",
+    "scipy>=1.6.0",
+    "tqdm>=4.55.0"
+]
+
+# optional dependencies for examples in pomdp_problems
+[project.optional-dependencies]
+problems = [
+    "matplotlib>=3.2.0",       # for some utilities
+    "pygame>=2.0.2",           # for some tests
+    "opencv-python>=4.5.2.0"   # for some tests
+]
+
+[project.urls]
+Home = "https://github.com/h2r/pomdp-py"
diff --git a/setup.py b/setup.py
index c5d6796a..5b01767d 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 
-from setuptools import setup, find_packages, Extension
+from setuptools import setup, Extension
 from Cython.Build import cythonize
 import os.path
 
@@ -26,30 +26,11 @@ def build_extensions(pkg_name, major_submodules):
                                            "algorithms",
                                            "utils",
                                            "representations.distribution",
-                                           "representations.belief"])
-extensions.extend(build_extensions("pomdp_problems", ["tiger.cythonize",
-                                                      "rocksample.cythonize"]))
+                                           "representations.belief",
+                                           "problems.tiger.cythonize",
+                                           "problems.rocksample.cythonize"])
 
-setup(name='pomdp-py',
-      packages=find_packages(),
-      version='1.3.3',
-      description='Python POMDP Library.',
-      long_description=long_description,
-      long_description_content_type="text/x-rst",
-      install_requires=[
-          'Cython==3.0.8',
-          'numpy',
-          'scipy',
-          'tqdm',
-          'matplotlib',
-          'pygame',        # for some tests
-          'opencv-python',  # for some tests
-      ],
-      license="MIT",
-      author='Kaiyu Zheng',
-      author_email='kzheng10@cs.brown.edu',
-      keywords = ['Partially Observable Markov Decision Process', 'POMDP'],
-      ext_modules=cythonize(extensions,
+setup(ext_modules=cythonize(extensions,
                             build_dir="build",
                             compiler_directives={'language_level' : "3"}),
       package_data={"pomdp_py": ["*.pxd", "*.pyx", "*.so", "*.c"],
diff --git a/tests/test_conversion_pomdp-solve.py b/tests/test_conversion_pomdp-solve.py
index baeb30c2..399432d7 100644
--- a/tests/test_conversion_pomdp-solve.py
+++ b/tests/test_conversion_pomdp-solve.py
@@ -6,12 +6,17 @@
 import pomdp_py
 import subprocess
 import inspect
-from pomdp_py.utils.test_utils import *
 from pomdp_py.utils.interfaces.conversion import to_pomdp_file
+from pomdp_py.problems.tiger import make_tiger
 
 
 description="testing conversion to .pomdp file"
 
+def remove_files(pattern):
+    file_list = glob.glob(pattern)
+    for file_path in file_list:
+        os.remove(file_path)
+
 def test_pomdp_file_conversion(pomdp_solve_path):
     print("[testing] test_pomdp_file_conversion")
     tiger = make_tiger()
diff --git a/tests/test_conversion_pomdpx.py b/tests/test_conversion_pomdpx.py
index 6850c23e..958743ce 100644
--- a/tests/test_conversion_pomdpx.py
+++ b/tests/test_conversion_pomdpx.py
@@ -1,8 +1,8 @@
 # Tests the conversion to .pomdpx file format code
 
 import sys
-from pomdp_py.utils.test_utils import *
 from pomdp_py.utils.interfaces.conversion import to_pomdpx_file
+from pomdp_py.problems.tiger import make_tiger
 import os
 
 description="testing conversion to .pomdpx file"
diff --git a/tests/test_sarsop.py b/tests/test_sarsop.py
index f95ed5ef..82d490ca 100644
--- a/tests/test_sarsop.py
+++ b/tests/test_sarsop.py
@@ -1,9 +1,9 @@
 import sys
 import pomdp_py
 import subprocess
-from pomdp_py.utils.test_utils import *
 from pomdp_py.utils.interfaces.conversion import to_pomdp_file
 from pomdp_py.utils.interfaces.solvers import sarsop
+from pomdp_py.problems.tiger import make_tiger
 import os
 import io
 
diff --git a/tests/test_tree_debugger.py b/tests/test_tree_debugger.py
index 2de2a0a7..7f9ebab7 100644
--- a/tests/test_tree_debugger.py
+++ b/tests/test_tree_debugger.py
@@ -1,5 +1,5 @@
 import random
-from pomdp_problems.tiger import TigerProblem, test_planner
+from pomdp_py.problems.tiger import TigerProblem, test_planner
 import pomdp_py
 from pomdp_py.utils.debugging import TreeDebugger
 
diff --git a/tests/test_vi_pruning.py b/tests/test_vi_pruning.py
index 7fc35ddf..02810a90 100644
--- a/tests/test_vi_pruning.py
+++ b/tests/test_vi_pruning.py
@@ -1,9 +1,10 @@
 import sys
 import pomdp_py
 import subprocess
-from pomdp_py.utils.test_utils import *
 from pomdp_py.utils.interfaces.conversion import to_pomdp_file
 from pomdp_py.utils.interfaces.solvers import vi_pruning
+from pomdp_py.problems.tiger import make_tiger
+
 import os
 import io