Wrapping custom gym with PyEnvironment TF-Agents

Question

I've got a custom gym environment which has a render method I can call with go_env.render(mode="human") (which draws a pyglet canvas). The gym I've got works with go_env = gym.make('gym_go:go-v0', size=args.boardsize, komi=args.komi) . How would I write the corresponding environment wrapper in TF-Agents? Currently I have: env = suite_gym.load('gym_go:go-v0', gym_kwargs={'size':3,'komi':0}, render_kwargs={'mode':'terminal'}) . But this throws a TypeError as render() on PyEnvironment has its own mode:

render(
    mode: Text = 'rgb_array'
) -> Optional[types.NestedArray]

How do I wrap the gym environment correctly with my parameters, and make sure it renders properly?

Answer 1

Turns out you don't need to pass it in renderkwargs , you can pass the rendering mode directly into the wrapped class like so:

env = suite_gym.load('gym_go:go-v0', gym_kwargs={'size':3,'komi':0})
env.render('terminal')

This works with custom modes too, as long as you override the render method in your custom gym implementation.