torch.distributed.rpc should be able to help. Here is a list of tutorials.
The use case looks similar to the following two examples:

https://pytorch.org/tutorials/intermediate/rpc_tutorial.html#distributed-reinforcement-learning-using-rpc-and-rref
https://pytorch.org/tutorials/intermediate/rpc_async_execution.html#batch-processing-cartpole-solver 1