diff --git a/README.md b/README.md
index 6625228..f20a4a5 100644
--- a/README.md
+++ b/README.md
@@ -10,10 +10,11 @@ The original MATLAB RTD Planner and ARMOUR ported into Python
 2. Install *virtualenv* with `py -m pip install virtualenv`
 3. Create a new Virtual Environment with `py -m venv rtd_venv`
 4. Activate the Virtual Environment with `rtd_venv\scripts\activate`
-5. Install *rtd-code-python* and its dependencies with `py -m pip install -r requirements.txt`
-6. Optionally, open a terminal inside `scripts\demos` and run the demo scripts
-
+5. Install *rtd-code-python* and its dependencies with `py -m pip install -r requirements.txt && py -m pip install -e .`
+6. Separately install [*zonopy*](https://github.com/roahmlab/zonopy) and [*zonopy-robots*](https://github.com/roahmlab/zonopy-robots).
+7. Optionally, open a terminal inside `scripts\demos` and run the demo scripts (the `urdfs` folder can be found from the [*armour* repository](https://github.com/roahmlab/armour))
 
+**Note**: You may need to apply the fix in https://github.com/fishbotics/urchin/issues/9 to load some of the armour URDF's
 
 ## Key Features
 - ARMOUR (`armour`)
diff --git a/armour/ArmourGoal.py b/armour/ArmourGoal.py
index 68ee8a9..dd82949 100644
--- a/armour/ArmourGoal.py
+++ b/armour/ArmourGoal.py
@@ -79,7 +79,7 @@ def reset(self, **options):
     def create_plot_data(self, time: float = None) -> list[Actor]:
         # generate mesh
         config = self.goal_position
-        fk: OrderedDict[Trimesh, Matnp] = self.arm_agent.info.robot.visual_trimesh_fk(cfg=config)
+        fk: OrderedDict[Trimesh, Matnp] = self.arm_agent.info.urdf.visual_trimesh_fk(cfg=config)
         meshes = [mesh.copy().apply_transform(transform) for mesh, transform in fk.items()]
         
         self.plot_data: list[Actor] = list()
diff --git a/armour/ArmourPlanner.py b/armour/ArmourPlanner.py
index 679e141..5ce895c 100644
--- a/armour/ArmourPlanner.py
+++ b/armour/ArmourPlanner.py
@@ -6,9 +6,9 @@
 from rtd.sim.world import WorldState
 from armour.reachsets import JRSGenerator, FOGenerator, IRSGenerator, JLSGenerator
 from armour.trajectory import ArmTrajectoryFactory
-from zonopy.robots2.robot import ZonoArmRobot
+from zonopyrobots import ZonoArmRobot
 from urchin import URDF
-
+from armour.agent import ArmourAgentInfo
 
 
 class ArmourPlanner(RtdPlanner, Options):
@@ -22,18 +22,18 @@ def defaultoptions() -> dict:
         }
         
         
-    def __init__(self, trajOptProps: TrajOptProps, robot: URDF, params: ZonoArmRobot, **options):
+    def __init__(self, trajOptProps: TrajOptProps, robot: ArmourAgentInfo, **options):
         # initialize base classes
         RtdPlanner.__init__(self)
         Options.__init__(self)
         # initialize using given options
         self.mergeoptions(options)
         self.rsGenerators = dict()
-        self.rsGenerators["jrs"] = JRSGenerator(params, traj_type=options["traj_type"])
-        self.rsGenerators["fo"] = FOGenerator(params, self.rsGenerators["jrs"], smooth_obs=options["smooth_obs"])
+        self.rsGenerators["jrs"] = JRSGenerator(robot, traj_type=options["traj_type"])
+        self.rsGenerators["fo"] = FOGenerator(robot, self.rsGenerators["jrs"], smooth_obs=options["smooth_obs"])
         if options["input_constraints_flag"]:
-            self.rsGenerators["irs"] = IRSGenerator(params, self.rsGenerators["jrs"], use_robost_input=options["use_robust_input"])
-            self.rsGenerators["jls"] = JLSGenerator(params, self.rsGenerators["jrs"])
+            self.rsGenerators["irs"] = IRSGenerator(robot, self.rsGenerators["jrs"], use_robost_input=options["use_robust_input"])
+            self.rsGenerators["jls"] = JLSGenerator(robot, self.rsGenerators["jrs"])
         
         # create the trajectory factory
         self.trajectoryFactory = ArmTrajectoryFactory(trajOptProps, options["traj_type"])
diff --git a/armour/agent/ArmKinematics.py b/armour/agent/ArmKinematics.py
index 4d384dd..80e04a3 100644
--- a/armour/agent/ArmKinematics.py
+++ b/armour/agent/ArmKinematics.py
@@ -46,7 +46,7 @@ def get_link_transform(self, time):
     
     
     def get_link_transform_from_config(self, config):
-        fk: OrderedDict[Trimesh, NDArray] = self.arm_info.robot.visual_trimesh_fk(cfg=config)
+        fk: OrderedDict[Trimesh, NDArray] = self.arm_info.urdf.visual_trimesh_fk(cfg=config)
         return fk.values()[1:]
     
     
@@ -117,6 +117,7 @@ def get_link_rotations_and_translations(self, time_or_config: float | NDArray =
                     if dim == 3:
                         # rotation matrix of current link
                         axis_pred = self.arm_info.joints[idx].axes
+                        # TODO
                         R_succ = R_pred*self.arm_info.robot.Bodies[idx].Joint.JointToParentTransform[:3, :3]*axang2rotm(axis_pred, joint_rot)
                     else:
                         # rotation matrix of current link
@@ -132,6 +133,7 @@ def get_link_rotations_and_translations(self, time_or_config: float | NDArray =
                 
                 case 'fixed':
                     if dim == 3:
+                        # TODO
                         R_succ = R_pred*self.arm_info.robot.Bodies[idx].Joint.JointToParentTransform[:3, :3]
                     else:
                         # rotation matrix of current link assumed same as predecessor
diff --git a/armour/agent/ArmourAgentCollision.py b/armour/agent/ArmourAgentCollision.py
index fb25e44..c7b8730 100644
--- a/armour/agent/ArmourAgentCollision.py
+++ b/armour/agent/ArmourAgentCollision.py
@@ -36,7 +36,7 @@ def getCollisionObject(self, q: Matnp = None, time: float = None) -> CollisionOb
         elif time is not None and q is None:
             config = self.arm_state.get_state(np.array([time])).position   # position at given time
             
-        fk: OrderedDict[Trimesh, Matnp] = self.arm_info.robot.collision_trimesh_fk(cfg=config)
+        fk: OrderedDict[Trimesh, Matnp] = self.arm_info.urdf.collision_trimesh_fk(cfg=config)
         meshes = [mesh.copy().apply_transform(transform) for mesh, transform in fk.items()]
         return CollisionObject(meshes, id(self.arm_info))
     
diff --git a/armour/agent/ArmourAgentInfo.py b/armour/agent/ArmourAgentInfo.py
index 6922e5c..3be9b1e 100644
--- a/armour/agent/ArmourAgentInfo.py
+++ b/armour/agent/ArmourAgentInfo.py
@@ -1,7 +1,8 @@
 from rtd.entity.components import BaseInfoComponent
 from rtd.util.mixins import Options
-from zonopy.robots2.robot import ZonoArmRobot
+from zonopyrobots import ZonoArmRobot
 from urchin import URDF
+import torch
 
 # define top level module logger
 import logging
@@ -19,20 +20,23 @@ def defaultoptions() -> dict:
         '''
         return {
             "M_min_eigenvalue": 0.002,
-            "gravity": [0, 0, -9,81],
+            "gravity": [0, 0, -9.81],
             "transmission_inertia": None,
             "buffer_dist": 0,
+            "torch_device": torch.device("cpu"),
         }
     
     
-    def __init__(self, robot: URDF, params: ZonoArmRobot, **options):
+    def __init__(self, robot: URDF, **options):
          # initialize base classes
         BaseInfoComponent.__init__(self)
         Options.__init__(self)
         # initialize using given options
-        self.mergeoptions(options)
-        self.robot: URDF = robot.copy()
-        self.params = params
+        options = self.mergeoptions(options)
+
+        # Save a ZonoArmRobot object generated from this URDF
+        self.params = ZonoArmRobot.load(robot, device=options["torch_device"], dtype=torch.double)
+        self.urdf = self.params.urdf
         
         # initialzie
         self.reset()
@@ -45,7 +49,7 @@ def reset(self, **options):
         options = self.mergeoptions(options)        
         # fill in other dependent factors
         #self.n_links_and_joints = self.params.nomianal.num_joints
-        self.n_q = len(self.robot.actuated_joints)
+        self.n_q = self.params.dof
         #self.body_joint_index = self.params.nominal.q_index
         self.M_min_eigenvalue = options["M_min_eigenvalue"]
         self.gravity = options["gravity"]
diff --git a/armour/agent/ArmourAgentVisual.py b/armour/agent/ArmourAgentVisual.py
index ebd771d..06554a3 100644
--- a/armour/agent/ArmourAgentVisual.py
+++ b/armour/agent/ArmourAgentVisual.py
@@ -55,7 +55,7 @@ def create_plot_data(self, time: float = None) -> list[Actor]:
 
         # generate mesh
         config = self.arm_state.get_state(np.array([time])).position
-        fk: OrderedDict[Trimesh, Matnp] = self.arm_info.robot.visual_trimesh_fk(cfg=config)
+        fk: OrderedDict[Trimesh, Matnp] = self.arm_info.urdf.visual_trimesh_fk(cfg=config)
         meshes = [mesh.copy().apply_transform(transform) for mesh, transform in fk.items()]
         
         self.plot_data: list[Actor] = list()
@@ -86,7 +86,7 @@ def plot(self, time: float = None):
 
         # generate mesh
         config = self.arm_state.get_state(np.array([time])).position
-        fk: OrderedDict[Trimesh, Matnp] = self.arm_info.robot.visual_trimesh_fk(cfg=config)
+        fk: OrderedDict[Trimesh, Matnp] = self.arm_info.urdf.visual_trimesh_fk(cfg=config)
         meshes = [mesh.copy().apply_transform(transform) for mesh, transform in fk.items()]
 
         for actor, mesh in zip(self.plot_data, meshes):
diff --git a/armour/reachsets/FOGenerator.py b/armour/reachsets/FOGenerator.py
index 11d8c8d..40a3f75 100644
--- a/armour/reachsets/FOGenerator.py
+++ b/armour/reachsets/FOGenerator.py
@@ -1,9 +1,9 @@
 from rtd.entity.states import EntityState
 from rtd.planner.reachsets import ReachSetGenerator
 from armour.reachsets import FOInstance, JRSGenerator
-from zonopy.conSet.polynomial_zonotope.poly_zono import polyZonotope
+from zonopy import polyZonotope
 from itertools import combinations
-from zonopy.kinematics import FO as FOcc
+# from zonopy.kinematics import FO as FOcc
 import torch
 import numpy as np
 
@@ -26,7 +26,7 @@ def __init__(self, robot, jrsGenerator: JRSGenerator, smooth_obs: bool = False,
         if obs_frs_combs is None:
             obs_frs_combs = {'maxcombs': 200, 'combs': None}
         self.cache_max_size = 0 # we don't want to cache any FO
-        self.robot = robot
+        self.robot = robot.params
         self.jrsGenerator: JRSGenerator = jrsGenerator
         self.obs_frs_combs: dict = obs_frs_combs
         if self.obs_frs_combs['combs'] is None:
@@ -41,7 +41,7 @@ def generateReachableSet(self, robotState: EntityState) -> dict[int, FOInstance]
         '''
         jrsInstance = self.jrsGenerator.getReachableSet(robotState, ignore_cache=True)
         logger.info("Generating forward occupancy!")
-        forwardocc = list(FOcc.forward_occupancy(jrsInstance[1].R, self.robot)[0].values())[1:jrsInstance[1].n_q+1]
+        forwardocc = list(forward_occupancy(jrsInstance[1].R, self.robot)[0].values())[1:jrsInstance[1].n_q+1]
         return {1: FOInstance(forwardocc, jrsInstance[1], self.obs_frs_combs)}
                     
     
@@ -52,4 +52,39 @@ def generate_combinations_upto(n: int) -> list[list[int]]:
         '''
         combs = [list(combinations(np.arange(i+1), 2)) for i in range(0, n)]
         combs[0].append(0)
-        return combs
\ No newline at end of file
+        return combs
+
+# Pulled from zonopy-ext
+from zonopy import polyZonotope, matPolyZonotope, batchPolyZonotope, batchMatPolyZonotope
+from collections import OrderedDict
+from zonopyrobots.kinematics import forward_kinematics
+from zonopyrobots import ZonoArmRobot
+
+from typing import Union, Dict, List, Tuple
+from typing import OrderedDict as OrderedDictType
+
+# Use forward kinematics to get the forward occupancy
+# Note: zono_order=2 is 5 times faster than zono_order=20 on cpu
+def forward_occupancy(rotatotopes: Union[Dict[str, Union[matPolyZonotope, batchMatPolyZonotope]],
+                                         List[Union[matPolyZonotope, batchMatPolyZonotope]]],
+                      robot: ZonoArmRobot,
+                      zono_order: int = 20,
+                      links: List[str] = None,
+                      link_zono_override: Dict[str, polyZonotope] = None,
+                      ) -> Tuple[OrderedDictType[str, Union[polyZonotope, batchPolyZonotope]],
+                                 OrderedDictType[str, Union[Tuple[polyZonotope, matPolyZonotope],
+                                                            Tuple[batchPolyZonotope, batchMatPolyZonotope]]]]:
+    
+    link_fk_dict = forward_kinematics(rotatotopes, robot, zono_order, links=links)
+    urdf = robot.urdf
+    link_zonos = {name: robot.link_data[urdf._link_map[name]].bounding_pz for name in link_fk_dict.keys()}
+    if link_zono_override is not None:
+        link_zonos.update(link_zono_override)
+    
+    fo = OrderedDict()
+    for name, (pos, rot) in link_fk_dict.items():
+        link_zono = link_zonos[name]
+        fo_link = pos + rot@link_zono
+        fo[name] = fo_link.reduce_indep(zono_order)
+    
+    return fo, link_fk_dict
\ No newline at end of file
diff --git a/armour/reachsets/FOInstance.py b/armour/reachsets/FOInstance.py
index f3271c5..7ac2b0e 100644
--- a/armour/reachsets/FOInstance.py
+++ b/armour/reachsets/FOInstance.py
@@ -75,7 +75,7 @@ def genNLConstraint(self, worldState: WorldState) -> Callable:
                 grad_obs_constraint_pz_slice = lambda k: obs_constraint_pz.cpu().grad_center_slice_all_dep(k)
                 
                 # save
-                obs_constraints.append(lambda k: self.batch_obs_constraint(obs_constraint_pz_slice, grad_obs_constraint_pz_slice, torch.from_numpy(k).float()))
+                obs_constraints.append(lambda k: self.batch_obs_constraint(obs_constraint_pz_slice, grad_obs_constraint_pz_slice, k))
         
         # create the constraint callback
         return lambda k: self.eval_constraints(k, len(obs_constraints), obs_constraints, batch_size=100)
@@ -98,9 +98,10 @@ def batch_obs_constraint(self, c: Callable, grad_c: Callable, k) -> tuple[float,
     def eval_constraints(k, n_c: int, obs_constraints: list[Callable], batch_size: int = 1):
         h = np.zeros(n_c*batch_size)
         grad_h = np.zeros((n_c*batch_size, k.size))
+        k_tensor = torch.as_tensor(k)
         
         for i in range(n_c):
-            h[i*batch_size:(i+1)*batch_size], grad_h[i*batch_size:(i+1)*batch_size,:] = obs_constraints[i](k)
+            h[i*batch_size:(i+1)*batch_size], grad_h[i*batch_size:(i+1)*batch_size,:] = obs_constraints[i](k_tensor)
         
         grad_heq = None
         heq = None
diff --git a/armour/reachsets/IRSGenerator.py b/armour/reachsets/IRSGenerator.py
index 96b75b4..3f65d09 100644
--- a/armour/reachsets/IRSGenerator.py
+++ b/armour/reachsets/IRSGenerator.py
@@ -1,8 +1,8 @@
 from rtd.entity.states import EntityState
 from rtd.planner.reachsets import ReachSetGenerator
 from armour.reachsets import IRSInstance, JRSGenerator
-from zonopy.conSet.polynomial_zonotope.poly_zono import polyZonotope
-from zonopy.dynamics.RNEA import pzrnea
+from zonopy import polyZonotope
+from zonopyrobots.dynamics.RNEA import pzrnea
 import torch
 import numpy as np
 
diff --git a/armour/reachsets/IRSInstance.py b/armour/reachsets/IRSInstance.py
index cd851dd..bd38255 100644
--- a/armour/reachsets/IRSInstance.py
+++ b/armour/reachsets/IRSInstance.py
@@ -2,7 +2,7 @@
 from rtd.planner.reachsets import ReachSetInstance
 from rtd.sim.world import WorldState
 from armour.reachsets import JRSInstance
-from zonopy.conSet.polynomial_zonotope.poly_zono import polyZonotope
+from zonopy import polyZonotope
 import numpy as np
 from rtd.util.mixins.Typings import Boundsnp
 
diff --git a/armour/reachsets/JLSGenerator.py b/armour/reachsets/JLSGenerator.py
index 96393d0..737c685 100644
--- a/armour/reachsets/JLSGenerator.py
+++ b/armour/reachsets/JLSGenerator.py
@@ -1,7 +1,7 @@
 from rtd.entity.states import EntityState
 from rtd.planner.reachsets import ReachSetGenerator
 from armour.reachsets import JLSInstance, JRSGenerator
-from zonopy.conSet.polynomial_zonotope.poly_zono import polyZonotope
+from zonopy import polyZonotope
 import torch
 import numpy as np
 
diff --git a/armour/reachsets/JLSInstance.py b/armour/reachsets/JLSInstance.py
index b619d85..584ce16 100644
--- a/armour/reachsets/JLSInstance.py
+++ b/armour/reachsets/JLSInstance.py
@@ -2,7 +2,7 @@
 from rtd.planner.reachsets import ReachSetInstance
 from rtd.sim.world import WorldState
 from armour.reachsets import JRSInstance
-from zonopy.conSet.polynomial_zonotope.poly_zono import polyZonotope
+from zonopy import polyZonotope
 import numpy as np
 from rtd.util.mixins.Typings import Boundsnp
 
diff --git a/armour/reachsets/JRSGenerator.py b/armour/reachsets/JRSGenerator.py
index cc634d2..904098a 100644
--- a/armour/reachsets/JRSGenerator.py
+++ b/armour/reachsets/JRSGenerator.py
@@ -1,8 +1,9 @@
 from rtd.entity.states import EntityState
 from rtd.planner.reachsets import ReachSetGenerator
 from armour.reachsets import JRSInstance
-from zonopy.joint_reachable_set.gen_jrs import JrsGenerator as ZonoJRSGenerator
-import zonopy.trajectories as zpt
+import zonopyrobots as zpr
+# from zonopy.joint_reachable_set.gen_jrs import JrsGenerator as ZonoJRSGenerator
+# import zonopy.trajectories as zpt
 import numpy as np
 
 # define top level module logger
@@ -28,8 +29,8 @@ def __init__(self, robot, taylor_degree: int = 1, add_ultimate_bound: bool = Tru
         self.add_ultimate_bound = add_ultimate_bound
         self.traj_type = traj_type
         # initialize zonopy's JRSGenerator
-        traj_class = zpt.PiecewiseArmTrajectory if traj_type=="piecewise" else zpt.BernsteinArmTrajectory
-        self._jrnsgen = ZonoJRSGenerator(robot, traj_class, k_r=None, # k_r=self.controller.k_r
+        traj_class = zpr.trajectory.PiecewiseArmTrajectory if traj_type=="piecewise" else zpr.trajectory.BernsteinArmTrajectory
+        self._jrnsgen = zpr.JrsGenerator(robot.params, traj_class, k_r=None, # k_r=self.controller.k_r
                                          ultimate_bound=None, # ultimate_bound=self.controller.ultimate_bound
                                          batched=True, unique_tid=False)
     
diff --git a/armour/trajectory/PieceWiseArmTrajectory.py b/armour/trajectory/PiecewiseArmTrajectory.py
similarity index 100%
rename from armour/trajectory/PieceWiseArmTrajectory.py
rename to armour/trajectory/PiecewiseArmTrajectory.py
diff --git a/requirements.txt b/requirements.txt
index b986436..206aa2e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -43,7 +43,7 @@ python-fcl==0.7.0.4
 pyvista==0.39.0
 PyYAML==6.0
 requests==2.31.0
--e git+https://github.com/roahmlab/rtd-code-python.git@8c11653276d0edb203bb39cb490977f429d2a698#egg=rtd_code_python
+# -e git+https://github.com/roahmlab/rtd-code-python.git@8c11653276d0edb203bb39cb490977f429d2a698#egg=rtd_code_python
 scipy==1.10.1
 scooby==0.7.2
 six==1.16.0
@@ -64,4 +64,5 @@ typing_extensions==4.7.1
 urchin==0.0.27
 urllib3==2.0.2
 vtk==9.2.6
--e git+https://github.com/BuildingAtom/zonopy.git@5a7060ac303393d6bd2cb120d2f4002b88ee4b2e#egg=zonopy
+# -e git+https://github.com/BuildingAtom/zonopy.git@5a7060ac303393d6bd2cb120d2f4002b88ee4b2e#egg=zonopy
+websockets==12.0
diff --git a/scripts/demos/planner_demo.py b/scripts/demos/planner_demo.py
index 45fc249..c1008fc 100644
--- a/scripts/demos/planner_demo.py
+++ b/scripts/demos/planner_demo.py
@@ -7,7 +7,7 @@
     from rtd.planner.trajopt import TrajOptProps
     from rtd.sim.world import WorldState
     from rtd.sim.sensors import zonotope_sensor
-    from zonopy.robots2.robot import ZonoArmRobot
+    from zonopyrobots import ZonoArmRobot
     from urchin import URDF
     import os
     
@@ -18,7 +18,10 @@
     
     
     #-------------------- agent parameters --------------------#
+    import rtd
+    rtd_code_basepath = os.path.dirname(rtd.__file__)
     agent_urdf = "../../urdfs/kinova_arm/kinova_without_gripper.urdf"
+    agent_urdf = os.path.join(rtd_code_basepath, agent_urdf)
     add_uncertainty_to = 'all'  # choose 'all', 'link', 'none'
     links_with_uncertainty = ['dumbell_link']
     uncertain_mass_range = (0.97, 1.03)
@@ -41,8 +44,7 @@
     
     #-------------------- setup the agent --------------------#
     urdf_path = os.path.join(os.path.dirname(__file__), agent_urdf)
-    robot = URDF.load(urdf_path)
-    params = ZonoArmRobot.load(robot)
+    robot_urdf = URDF.load(urdf_path)
 
     vel_limits = [[-1.3963, -1.3963, -1.3963, -1.3963, -1.2218, -1.2218, -1.2218],
                   [1.3963,  1.3963,  1.3963,  1.3963,  1.2218,  1.2218,  1.2218]]
@@ -51,7 +53,7 @@
     transmision_inertia = [8.02999999999999936, 11.99620246153036440, 9.00254278617515169, 11.58064393167063599, 8.46650409179141228, 8.85370693737424297, 8.85873036646853151]
     M_min_eigenvalue = 5.095620491878957
     
-    agent_info = ArmourAgentInfo(robot, params, joint_velocity_limits=vel_limits, joint_torque_limits=input_limits,
+    agent_info = ArmourAgentInfo(robot_urdf, joint_velocity_limits=vel_limits, joint_torque_limits=input_limits,
                                  transmission_inertia=transmision_inertia, M_min_eigenvalue=M_min_eigenvalue)
     
     armour_controller = ArmourController
@@ -84,8 +86,7 @@
     
     planner = ArmourPlanner(
         trajOptProps=trajOptProp,
-        robot=sim.agent,
-        params=params,
+        robot=agent_info,
         input_constraints_flag=input_constraints_flag,
         use_robust_input=use_robust_input,
         smooth_obs=smooth_obs,