How can I use all cores of the CPU?

[masterkey2000]

Hello,

i use tf-agents to train a DQN. I wonder why only one core of my CPU is used to compute the environment. I have a Ryzen 9 5950x with 16 cores. I would like to use the CPU to its full capacity. What do I have to do for this?

Edit:

Hier is the relevant code I use:

def train_step(self, n_steps):
    
    env_steps = tf_metrics.EnvironmentSteps()
    #num_episodes = tf_metrics.NumberOfEpisodes()
    rew = TFSumOfRewards()    
    action_hist = tf_metrics.ChosenActionHistogram(
        name='ChosenActionHistogram', dtype=tf.int32, buffer_size=1000
        )
    
    #add reply buffer and metrict to the observer
    replay_observer = [self.replay_buffer.add_batch]   
    train_metrics = [env_steps, rew]
    
    self.replay_buffer.clear()
    
    driver = dynamic_episode_driver.DynamicEpisodeDriver(
        self.train_env, self.collect_policy, observers=replay_observer + train_metrics, num_episodes=self.collect_episodes)
    
    final_time_step, policy_state = driver.run()
    
    print('Number of Steps: ', env_steps.result().numpy())
    
    for train_metric in train_metrics:
        train_metric.tf_summaries(train_step=self.global_step, step_metrics=train_metrics)
    
    # Convert the replay buffer to a tf.data.Dataset 
    # Dataset generates trajectories with shape [Bx2x...]
    AUTOTUNE = tf.data.experimental.AUTOTUNE
    dataset = self.replay_buffer.as_dataset(
        num_parallel_calls=AUTOTUNE, 
        sample_batch_size=self.batch_size, 
        num_steps=(self.train_sequence_length + 1)).prefetch(AUTOTUNE)                      
   
    iterator = iter(dataset)
           
    train_loss = None
    for _ in range(n_steps):
        # Sample a batch of data from the buffer and update the agent's network.
        experience, unused_info = next(iterator)
        train_loss = self.agent.train(experience)
       

def train_agent(self, n_epoch):
    for i in range(n_epoch):
        self.train_step(int(self.replay_buffer.num_frames().numpy()/self.batch_size))
        if(self.IsAutoStoreCheckpoint == True):
            self.store_check_point()   
    pass

[billh0420]

I used the following:

`

### TFPyEnvironment
def create_tf_env():
    py_env = SimpleEnv()
    tf_env = tf_py_environment.TFPyEnvironment(py_env)
    return tf_env

def create_batch_tf_env():
    ### Important
    batch_size = 16 ### This works great

    ### Note: not sure about multithreading should be False or True (doesn't seem to affect execution time).
    ### Note: not sure about isolation should be False or True (doesn't seem to affect execution time).
    ### Note: If isolation is True, then multithreading must be False (according to documentation).
    multithreading = False
    isolation = True
    assert not (multithreading and isolation)

    py_envs = [SimpleEnv() for _ in range(0, batch_size)]
    tf_env = tf_agents.environments.batched_py_environment.BatchedPyEnvironment(envs=py_envs, multithreading=multithreading)
    tf_env = TFPyEnvironment(tf_env, isolation=isolation)
    return tf_env

`