Question about performance and the configuration of the environment

[PaulPruenster]

Hi,
I have a question about performance and the configuration of the environment. I have a Gym environment that inside uses rlbench to train a reinforcement learning model. When measuring the performance of the simulation, I get about 0.12s average per task.step(), so given I have 50 num_steps per step of the gym environment, one step takes about 6 seconds. This is obviously way to slow to train a RL model, so I was wondering if maybe we missed a configuration of some sort, it seems to me like the simulation is running in real time and not as fast as possible.

This is the setup:

 self.env = Environment(
            action_mode=MoveArmThenGripper(
                arm_action_mode=JointPosition(), gripper_action_mode=Discrete()),
            obs_config=obs_config,
            headless=True, shaped_rewards=True)
        
self.env.launch()
self.task = self.env.get_task(ReachTarget)

This is one step:

def step(self, Z):
        num_steps = 50
        sum_reward = 0.0
        time = 0.0
        print(Z)
        for i in range(num_steps):
            action = self.decoder.decode_from_latent_variable(torch.from_numpy(Z), time)

            te = time.time()
            obs, reward, _ = self.task.step(action.detach().numpy()) # <=== SLOW HERE
            print(time.time() - te) # Output: 0.12534880638122559
            
            sum_reward += reward
            time += 1 / num_steps

        print(sum_reward)
        self.rewards.append(sum_reward)
        state = self.state_from_obs(obs)
        return state, sum_reward, True, True, {}

These are some versions we are using:
rlbench @ git+https://github.com/stepjam/RLBench.git@b80e51feb3694d9959cb8c0408cd385001b01382
CoppeliaSim_Edu_V4_1_0_Ubuntu20_04
Python 3.12.5