Distributed training with Kubernetes

[jlewi]

Opening this issue to start a discussion about whether it would be worth investing to make it easy to run tensorflow agents K8s.

For some inspiration you can look at TfJob CRD.

Some questions:

1. Is there a need to be able to distribute the environments across multiple machines?
2. What is the communication pattern between the simulations and TensorFlow job?
   * Is data fetched from all simulations simultaneously?
   * Does each simulation need to be individually addressable?

[danijar]

This is an interesting topic. I would imagine that a simple distribution pattern would be to run multiple instances of the current code on multiple machine, each simulating their own set of environments. The network parameters would need to be shared using parameter servers.

[jlewi]

@danijar So basically you want to do data parallel replicated training with TensorFlow
each of which is getting data from a set of environments running colocated with TF.

I think that would work out of the box with the TfJob CRD. The CRD will take care of provisioning a set of workers and parameter servers. The workers can run any code (docker image); so they could use tensorflow/agents to launch a bunch of environments as subprocesses.

You could also move the environments into 1 or more sidecars and communicate via the filesystem or sockets. This might be beneficial if

• You want better monitoring/isolation of the simulation environments compared to TF workers
• You have different docker images for your TensorFlow process and environments

Would it be beneficial to run the simulations on different machines from the workers?

It sounds like the pattern is many simulations to 1 worker; i.e. you have N workers and each worker uses B simulations to generate a batch of size B.

If tensorflow/agents adds an RPC mode of communication, then I think K8s should be able to handle distributing both TensorFlow workers and the simulations across multiple machines.

Using StatefulSets you could start N * B simulations; each with a stable DNS name. Each worker could then use B of these simulations.

[danijar]

I think in many scenarios it makes sense to simulate and train on the same machine, and just scale the number of those machines. That's mainly because it seems to reduce communication overhead significantly by only sharing parameters, not all the collected data. I'd be super curious to know if that works out of the box as you mentioned it might!

In some new scenarios we cannot simulate the environments on the same machine, let's say because the data is generated by some specialized service or on a real robot. In that case, we could implement a Gym environment that under the hood communicates with the data service via RPCs. But instead of communicating for every frame, we might want to send a policy to the data collector and get back full trajectories, which would require some more changes.

[jlewi]

I think in many scenarios it makes sense to simulate and train on the same machine, and just scale the number of those machines. That's mainly because it seems to reduce communication overhead significantly by only sharing parameters, not all the collected data. I'd be super curious to know if that works out of the box as you mentioned it might!

So would I. Is there a simple example we could try to parallelize across multiple machines using TfJob?

[danijar]

Maybe something like this: https://gist.github.com/danijar/1dde5eb528b61560734c16d4fd0d93a1

We will probably have to modify agents/scripts/train.py though.

[jlewi]

@cwbeitel any interest in trying to turn @danijar's gist into an example running on K8s?

[cwbeitel]

Sure!

[cwbeitel]

Still working on it! Hit some issues with my kube deployment when trying to re-deploy with accelerators. I'll push a version here in a minute to share for discussion.

[cwbeitel]

So here's the work in progress, let me know if you guys have thoughts about this model for how model code is built into images and run. It doesn't have to be so prescriptive but this is a workable model for now.

[cwbeitel]

OK so here's a cleaner version of that PR kubeflow/trainer#159. So next up is extending the example with some form of parallelization initially with the environments still running in the same containers as where replicas of the model are being trained. @danijar @jlewi do you guys have thoughts on including the example in tf/k8s vs tf/agents? I don't want to go out of scope on k8s and this issue was originally posted here.

[danijar]

Cool, I looked at the PR and it looks quite nice. The examples directory of the k8s repo is probably a good place for this code. I'm planning on adding some tutorials for Agents and would be happy to include a page about this and a link to the example.

Did you do a full distributed training run and have some graphs to look at?

[cwbeitel]

Thank you. Also made the example a bit simpler per @jlewi's suggestion, see cwbeitel/k8s@790dfb1.

That sounds good about the tutorial happy to help user test those.

You should be able to tensorboard the logs from one of the runs (let me know if you can't access it, marked "make public")

tensorboard --logdir gs://dev01-181118-181500-k8s/jobs/tensorflow-20171117102413/20171117T182424-pybullet_ant

Rendering is running so I'll upload a gif to the readme when that's done / I have a chance. Have a suggestion of which graphs to include in the readme or did you just mean for browsing?

Haven't made it distributed yet, that's next up!

[danijar]

Thanks. TensorBoard is getting a 401 Unauthorized return code. Yep, checked the code and saw that it's single-instance training for now. Let me know if you run into problems when working on the distributed version, maybe I can help. For the readme, the graph with the mean score is probably most interesting, or a screenshot of all the scores on TensorBoard.

[cwbeitel]

Ok gs://agents-k8s-share/jobs/... should work going forward including gs://agents-k8s-share/jobs/tensorflow-20171117102413/20171117T182424-pybullet_ant

Figures:

Thanks happy to have help. Getting up to speed a bit on this class of algorithm and away this week. Maybe at some point we can write up in pseudocode what exactly we want to parallelize. Here's pseudocode for the KL and clipped versions of PPO:

## PPO-1 (KL)
for iteration do
  Run policy for T timesteps over N trajectories
  Estimate advantage function at all timesteps
  Do SGD on above objective for some number of epochs
  If KL too high, increase beta
  If KL too low, decrease beta
end for

## PPO-2 (Clipped objective)
for iteration do
  Run policy for T timesteps or N trajectories
  Estimate advantage function at all timesteps
  Do SGD on L^CLIP(theta) objective for some number of epochs
end for

It seems like we basically want to do distributed SGD in the region of some policy parameterization where individual agents accumulate their own experience trajectories, perform a bit of private SGD, then all share their results with a central policy vector? Or were you thinking more that there would be a single node for policy optimization but distributed accumulation of experience (i.e. many parallel actors but a single critic?)