_thread.lock object cannot be pickled when using Pytorch multiprocessing package with spawn method

Muye_Jia: <p>Hi!</p>
<p>I’m trying to start a multiprocessing task using PPO algorithm, it worked well when I was using TD3 algorithm but somehow it fails for PPO for the problem of _thread.lock() cannot be pickled UPON STARTING, i.e. it fails when the start() method is called.</p>
<pre><code class="lang-auto">import gymnasium as gym
import numpy as np
from AsyncPPO import Worker
from PPO_torch import PPO
import torch.multiprocessing as mp


NUM_WORKERS = 3
UPDATE_ITER = 3*32
MAX_STEPS = 10000
GAMMA = 0.99
TAU = 0.0002
EXP_EPS = 128*10


if __name__ == "__main__":
    env = gym.make('MountainCarContinuous-v0')
    obs_shape = env.observation_space.shape[0]
    prev_obs = np.zeros(obs_shape)
    action_shape = env.action_space.shape[0]

    global_agent = PPO(num_states=obs_shape, num_actions=action_shape, mp=True, critic_lr=0.0002, actor_lr=0.0002)
    global_agent.share_memory()

    success_counter = 0
    step_counter = 0

    workers = []
    mp.set_start_method("spawn")
    for i in range(NUM_WORKERS):
        worker = Worker(globalAC=global_agent, name=str(i))
        t = mp.Process(target=worker.work)
        t.start()

        workers.append(t)

    for thread in workers:
        if not thread.is_alive():
            thread.join()
</code></pre>
<p>And it gives the following error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/itca-control/rl_ws/mpc_ws/src/rl_network/gym_test/AsyncPPO_mountain_car.py", line 40, in &lt;module&gt;
    t.start()
  File "/usr/lib/python3.8/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/usr/lib/python3.8/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/usr/lib/python3.8/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/usr/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/usr/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/usr/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/usr/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
TypeError: cannot pickle '_thread.lock' object
[W CudaIPCTypes.cpp:15] Producer process has been terminated before all shared CUDA tensors released. See Note [Sharing CUDA tensors]
</code></pre>
<p>I’m using Pytorch 2.0.1+cu117, Python 3.8<br>
The worker class is as follow if it helps</p>
<pre><code class="lang-auto">class Worker(mp.Process):
    def __init__(self, globalAC: PPO, name: str,
                 env=None, render=None) -&gt; None:
        """
        Initialize worker trainer
        """
        if env is not None:
            self.env = gym.make(env)
            if render is not None:
                gym.make(env, render_mode=render)

        else:
            self.env = gym.make('MountainCarContinuous-v0', render_mode="human")

        self.num_states = self.env.observation_space.shape[0]
        self.num_actions = self.env.action_space.shape[0]

        self.global_model = globalAC
        self.local_model = PPO(self.num_states, self.num_actions, critic_lr=0.0003, actor_lr=0.0002)
        self.name = name
        self.update_iter = self.local_model.buffer.buffer_capacity

    def work(self, max_steps=10, alow=-1, aup=1, exp_ep=1280):
        """
        Let the worker interact with the env and collect experience.

        The worker collects experience and adds to the master's replay buffer.
        """
        total_step = 0
        success_count = 0

        for ep in range(max_steps):
            obs, _ = self.env.reset()
            done = False
            ep_reward = 0.0

            while not done:
                if total_step &gt; exp_ep:
                    action = self.local_model.sample_action(obs, a_lower=alow, a_upper=aup)
                    action = utils.to_numpy(action)
                else:
                    action = self.env.action_space.sample()

                obs, reward, done, _, _ = self.env.step(action)
                next_obs = np.squeeze(obs)
                self.local_model.buffer.record((obs, action, reward, next_obs, done))

                total_step += 1
                ep_reward += reward
                if total_step % self.update_iter == 0:
                    global_buffer = self.local_model.buffer.on_policy_collect()
                    self.global_model.update_from_buffer(global_buffer)

                    # Pull the update to local
                    self.local_model.actor.load_state_dict(self.global_model.actor.state_dict())
                    self.local_model.critic.load_state_dict(self.global_model.critic.state_dict())

                print("worker " + self.name + \
                      " Ep " + str(ep) + \
                      " Step " + str(total_step) + " reward " + str(200*reward))

                if reward &gt;= 100:
                    done = True

            if done:
                print("Worker " + self.name + " training successful")
                break
</code></pre>

Muye_Jia: <p>It turns out this problem occurs because I imported tensorboard and created SummaryWriter which for some reason messed up the thread lock object</p>

cangozpi: <aside class="quote no-group" data-username="Muye_Jia" data-post="2" data-topic="184953">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/muye_jia/48/62322_2.png" class="avatar"> Muye_Jia:</div>
<blockquote>
<p>ause I imported tensorboard and created Su</p>
</blockquote>
</aside>
<p>Hi, I am facing the same issue with tensorboard. Did you find any solution to this problem ?</p>

tanyjnaaman: <p>In my case, I needed a <code>SummaryWriter</code> to log stuff. But it has a file handler, which makes parallelism impossible - see <a href="https://github.com/pytorch/pytorch/issues/32046#issuecomment-573151743" class="inline-onebox" rel="noopener nofollow ugc">torch.save causes TypeError: can't pickle SwigPyObject objects · Issue #32046 · pytorch/pytorch · GitHub</a>.</p>
<p>A simple fix I used was to instantiate a <code>SummaryWriter</code> everytime I logged, rather than have a single instance.</p>
<p>i.e. instead of</p>
<pre><code class="lang-auto">self.writer = SummaryWriter(...)

# somewhere in the code
self.writer.log(...)
</code></pre>
<p>I did:</p>
<pre><code class="lang-auto">self.log_dir = ...

# somewhere in the code
SummaryWriter(self.log_dir).log(...)
</code></pre>