fix wp.index error and minor bugfixes

src/curobo/rollout/rollout_base.py

@@ -23,7 +23,7 @@ import torch
 # CuRobo
 from curobo.types.base import TensorDeviceType
 from curobo.types.math import Pose
-from curobo.types.robot import CSpaceConfig, State
+from curobo.types.robot import CSpaceConfig, State, JointState
 from curobo.types.tensor import (
     T_BDOF,
     T_DOF,
@@ -456,7 +456,7 @@ class RolloutBase:
         return out_metrics
 
     def get_metrics_cuda_graph(self, state: State):
-        return get_metrics(state)
+        return self.get_metrics(state)
 
     def rollout_fn(self, act):
         pass