diff --git a/kscale/formats/mjcf.py b/kscale/formats/mjcf.py
index 6b0551e..c4bce4f 100644
--- a/kscale/formats/mjcf.py
+++ b/kscale/formats/mjcf.py
@@ -1,4 +1,3 @@
-# mypy: disable-error-code="operator,union-attr"
 """Defines common types and functions for exporting MJCF files.
 
 API reference:
@@ -443,14 +442,14 @@ def __init__(
             compiler (Compiler, optional): The compiler settings.
         """
         self.robot_name = robot_name
-        self.output_dir = output_dir
+        self.output_dir = Path(output_dir)
         self.compiler = compiler
         self._set_clean_up()
         self.tree = ET.parse(self.output_dir / f"{self.robot_name}.xml")
 
     def _set_clean_up(self) -> None:
         try:
-            import mujoco
+            import mujoco  # type: ignore[import-not-found]
         except ImportError as e:
             raise ImportError(
                 "Please install the package with Mujoco as a dependency, using "
@@ -486,6 +485,9 @@ def adapt_world(self) -> None:
             compiler = self.compiler.to_xml(compiler)
 
         worldbody = root.find("worldbody")
+        if worldbody is None:
+            raise ValueError("No worldbody found in the XML file")
+
         new_root_body = Body(name="root", pos=(0, 0, 0), quat=(1, 0, 0, 0)).to_xml()
 
         # List to store items to be moved to the new root body
@@ -493,6 +495,7 @@ def adapt_world(self) -> None:
         # Gather all children (geoms and bodies) that need to be moved under the new root body
         for element in worldbody:
             items_to_move.append(element)
+
         # Move gathered elements to the new root body
         for item in items_to_move:
             worldbody.remove(item)
diff --git a/kscale/store/pybullet.py b/kscale/store/pybullet.py
index e9fe1ea..86e77a9 100644
--- a/kscale/store/pybullet.py
+++ b/kscale/store/pybullet.py
@@ -24,7 +24,7 @@ def main(args: Sequence[str] | None = None) -> None:
     parsed_args = parser.parse_args(args)
 
     try:
-        import pybullet as p
+        import pybullet as p  # type: ignore[import-not-found]
     except ImportError:
         raise ImportError("pybullet is required to run this script")
 
diff --git a/setup.py b/setup.py
index d93ca1f..03d0121 100644
--- a/setup.py
+++ b/setup.py
@@ -1,4 +1,4 @@
-# mypy: disable-error-code="import-untyped"
+# mypy: disable-error-code="import-untyped, import-not-found"
 #!/usr/bin/env python
 """Setup script for the project."""