Updating OmniKit when using OpenAI Gym and Isaac Sim


I am working on a reinforcement learning project that uses Open AI Gym to train a custom robot in Isaac SIM.

So far I have created my own sample inside omni.isaac.samples (isaac-sim/_build/linux-x86_64/release/exts/omni.isaac.samples/omni/isaac/samples/scripts). Here I have managed to import my robot, control it and make it navigate using a similar logic to the simple_navigation_preview example.

This is how my step function in my custom gym environment looks like:

def step(self, action):
        print("Number of steps ", self.numsteps)
        print("Action ", action)

        # Check if the action is within the action space
        assert self.action_space.contains(action), "%r (%s) invalid" % (action, type(action))

        left_wheel_rpm, right_wheel_rpm = cmd_vel_to_rpm(action[0], action[1])
        self._rc.control_command(left_wheel_rpm, right_wheel_rpm)

        angle_to_goal = self._rc.angle_to_goal()
        distance_to_target = self._rc.distance_to_target()
        extra_values = np.array([angle_to_goal, distance_to_target])

        lidar_scans = np.concatenate((self._rc.rear_lidar_depth, self._rc.front_lidar_depth))
        position = np.concatenate(([[distance_to_target]], [[angle_to_goal]]))
        self.current_state = np.concatenate((lidar_scans, position)).reshape(1, 50, 1)

        reward, done, info = self.calculate_reward(action)
        self.numsteps += 1

        if self.numsteps > self.max_episode_steps:
            done = True
            print("robot stepped 500 times")

        return self.current_state, reward, done, info

However, when the training starts, even though control commands are sent to the RobotController (self._rc)the editor is frozen and the robot does not move until the training stops - my guess is that the OmniKit editor is not synchronized with the gym environment python code since when the training stops, the robot reproduces all the actions that were supposed to be performed during training.

I saw in the jetracer reinforcement learning example (isaac-sim/python-examples/jetracer) that they are using the OmniKitHelper class to launch OmniKit from a Python environment and also update the OmniKit (self.omniverse_kit.update(self.dt)) whenever an action command is sent in the step function.

def step(self, action):
        print("Number of steps ", self.numsteps)
        transformed_action = self.transform_action(action)
        frame = 0
        total_reward = 0
        reward = 0
        while frame < 3:
            obs = self.jetracer.observations()

I would like to know how to implement an update function that can called whenever a new frame needs to be rendered (i.e. when the step function is executed). Is there a way to do this without using the OmniKitHelper and having to launch OmniKit from a Python environment? I would like to maintain my current solution where I launch my gym environment from the Omniverse Isaac Sim Editor menu. Thanks!