Problems running eval.py

[vyeevani]

Running the eval.py yields the following issues. Seems like a versioning problem somewhere in the docker config/jaxrl_m stuff

Traceback (most recent call last):
File "experiments/eval.py", line 304, in
app.run(main)
File "/home/robonet/myenv/lib/python3.8/site-packages/absl/app.py", line 308, in run
_run_main(main, args)
File "/home/robonet/myenv/lib/python3.8/site-packages/absl/app.py", line 254, in _run_main
sys.exit(main(argv))
File "experiments/eval.py", line 101, in main
agent, action_mean, action_std = load_checkpoint(
File "experiments/eval.py", line 74, in load_checkpoint
agent = agents[config["agent"]].create(
File "/home/robonet/code/bridge_data_v2/jaxrl_m/agents/continuous/gc_bc.py", line 164, in create
params = model_def.init(init_rng, actor=[(observations, goals)])["params"]
File "/home/robonet/code/bridge_data_v2/jaxrl_m/common/common.py", line 72, in call
out[key] = self.moduleskey
File "/home/robonet/code/bridge_data_v2/jaxrl_m/networks/actor_critic_nets.py", line 146, in call
distribution = distrax.MultivariateNormalDiag(
File "/home/robonet/myenv/lib/python3.8/site-packages/distrax/_src/distributions/mvn_diag.py", line 77, in init
loc = conversion.as_float_array(loc)
File "/home/robonet/myenv/lib/python3.8/site-packages/distrax/_src/utils/conversion.py", line 138, in as_float_array
if not isinstance(x, Array):
File "/usr/lib/python3.8/typing.py", line 769, in instancecheck
return self.subclasscheck(type(obj))
File "/usr/lib/python3.8/typing.py", line 777, in subclasscheck
raise TypeError("Subscripted generics cannot be used with"
TypeError: Subscripted generics cannot be used with class and instance checks

[vyeevani]

Distrax 0.1.3 seems to fix the problem. Consider moving the requirements to include this

[HomerW]

Ah yeah this is an unfortunate issue because for training we need distrax==0.1.2 (given the other dependencies). However, because the we have to use python 3.8 for evaluating on the robot (due to using ROS noetic), we get this error. The best solution I could come up with was to add a line in the Dockerfile that upgrades distrax for the environment used when evaluating on the robot.

https://github.com/rail-berkeley/bridge_data_v2/blob/main/Dockerfile#L11