Revert "[wingman -> rllib] IMPALA MultiDiscrete changes (#3967)"

[ericl]

What do these changes do?

There seems to be a bug here, isolated to the changes in vtrace.py. Reverting that file and patching up multi_from_logits to call from_logits fixes the regression. Unfortunately I don't have time to investigate this further so we should revert this change and re-land it once the bug is found.

Interestingly, vtrace.py still passes its original unit tests.

The regression is fairly straightforward to reproduce: on any GPU machine run:

atari-impala:
    env:
        grid_search:
            - BreakoutNoFrameskip-v4
    run: IMPALA
    config:
        sample_batch_size: 50
        train_batch_size: 500
        num_workers: 32
        num_envs_per_worker: 5
        clip_rewards: True
        lr_schedule: [
            [0, 0.0005],
            [20000000, 0.000000000001],
        ]

Within a few minutes it gets to 1m+ timesteps. Before the regression, you can expect 10+ reward at this point. After the regression, the reward will top out at 1-2.

Related issue number

Atari performance regression originally reported in #4329

[AmplabJenkins]

Test PASSed.
Refer to this link for build results (access rights to CI server needed):
https://amplab.cs.berkeley.edu/jenkins//job/Ray-Perf-Integration-PRB/36/
Test PASSed.

[ericl]

First time build. Skipping changelog.
[Ray-Perf-Integration-PRB] $ /bin/bash /tmp/hudson4942004597160488522.sh
PR 4332 does not contain ray-core, early exiting
Test PASSed.

@robertnishihara is this misdetection due to a revert or is it some other bug?

[pcmoritz]

@ericl This is the new performance testing jenkins, which only kicks in if the PR string contains ray-core afaik (@devin-petersohn), i.e. if something performance-related has changed. ray-core should probably be changed to core.

[ericl]

Ah ok, didn't realize it had a different name.

[hartikainen]

One potential source of the bug is if this guy has a wrong shape.

[ericl]

I checked the shapes here and it looked reasonable (at least they match).

[AmplabJenkins]

Test FAILed.
Refer to this link for build results (access rights to CI server needed):
https://amplab.cs.berkeley.edu/jenkins//job/Ray-PRB/12765/
Test FAILed.

[stefanpantic]

@ericl Hi Eric, I think I fixed the issue, I've opened a pull request here