Fix bug in shaped reward networks

[levmckinney]

Description

Addresses #543

Testing

• Add test for shaping actually being applied during predict and predict processed.
• Add test for ValueError being raised when you try to shape a wrapped reward net.

[codecov]

Codecov Report

Merging #544 (fc9db6b) into master (de36306) will increase coverage by 0.01%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #544      +/-   ##
==========================================
+ Coverage   96.95%   96.96%   +0.01%     
==========================================
  Files          84       84              
  Lines        7460     7486      +26     
==========================================
+ Hits         7233     7259      +26     
  Misses        227      227              

Impacted Files	Coverage Δ
src/imitation/rewards/reward_nets.py	97.14% <100.00%> (+0.06%)	⬆️
src/imitation/testing/reward_nets.py	100.00% <100.00%> (ø)
tests/rewards/test_reward_nets.py	100.00% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[AdamGleave]

I'd like to reduce code duplication but otherwise this looks fine. Please request re-review once addressed.

[dfilan]

FWIW I'm kind of confused re: what about ForwardWrapper makes it so that it primarily changes the behaviour of forward, and why it looks so different from the code for PredictProcessedWrapper.

[levmckinney]

FWIW I'm kind of confused re: what about ForwardWrapper makes it so that it primarily changes the behaviour of forward, and why it looks so different from the code for PredictProcessedWrapper.

Since the forward wrapper inherits from RewardNet without implementing forward, it still an abstract class. Thus, anyone subclassing ForwardWrapper will have to implement forward if they want to be able to create concrete instances of the class.

The reason it looks so different is that: for the ForwardWrapepr we don't want all the methods calls on instances of ForwardWrapper to be pass through to their self.base. On the other hand, for the PredictProcessedWrapper we generally do.

[AdamGleave]

LGTM.