Some fixes

omnigibson/prims/entity_prim.py

@@ -162,12 +162,14 @@ class EntityPrim(XFormPrim):
 
     def remove(self):
         # First remove all joints
-        for joint in self._joints.values():
-            joint.remove()
+        if self._joints is not None:
+            for joint in self._joints.values():
+                joint.remove()
 
         # Then links
-        for link in self._links.values():
-            link.remove()
+        if self._links is not None:
+            for link in self._links.values():
+                link.remove()
 
         # Finally, remove this prim
         super().remove()

tests/test_object_removal.py

@@ -44,6 +44,9 @@ def test_removal_and_readdition():
     )
     og.sim.import_object(apple2)
 
+    # Clear the stuff we added
+    og.sim.remove_object(apple2)
+
 @og_test
 def test_readdition():
     # Make a copy of NAMES
@@ -79,5 +82,5 @@ def test_readdition():
     # Check that NAMES has not changed
     assert NAMES == new_names
 
-if __name__ == "__main__":
-    test_removal_and_readdition()
+    # Clear the stuff we added
+    og.sim.remove_object(apple)