Internal change

PiperOrigin-RevId: 415563520
Change-Id: Ifebe12550ad63f04cc266854cfda3b4cf74f3ffa

brax/experimental/biggym/README.md

@@ -1,12 +1,47 @@
 # BIG-Gym
 
-BIG-Gym is a *crowd-sourcing* challenge for RL *environments* and *behaviors*, inspired by [BIG-Bench](https://github.com/google/BIG-bench). It is co-organized as part of NeurIPS 2021 [Ecology Theory for RL (EcoRL)](https://sites.google.com/view/ecorl2021/home) workshop. *Our goal is to create the "ImageNet" for continuous control, with diversity in agent morphologies, environment scenes, objects, and tasks.* We solicit submissions for two tracks: **Open-Ended Creativity Track** and **Goal-Oriented Competition Track**.
+BIG-Gym is a *crowd-sourcing* challenge for RL *environments* and *behaviors*, inspired by [BIG-Bench](https://github.com/google/BIG-bench). *Our goal is to create the "ImageNet" for continuous control, with diversity in agent morphologies, environment scenes, objects, and tasks.* We solicit submissions for two tracks: **Open-Ended Creativity Track** and **Goal-Oriented Competition Track**.
 
+```python
+from brax.experimental import biggym
 
-## Organizers
+# register all in registry/__init__.py
+biggym.register_all(verbose=True)
+
+# register a specific folder under registry/
+#   `biggym.ENVS_BY_TRACKS` shows which envs are registered under each track
+env_names, component_names, task_env_names = biggym.register(registry_name)
+
+# (optional) inspect and get default configurable parameters of an environment
+env_params, _ = biggym.inspect_env(env_names[0])
+
+# create an environment
+env = biggym.create(env_names[0], env_params=env_params)
+```
 
-Core organizers: [Shixiang Shane Gu](https://sites.google.com/view/gugurus/home), [Hiroki Furuta](https://frt03.github.io/), [Manfred Diaz](https://manfreddiaz.github.io/)
+Challenge details (timelines, submission instructions) are [here](https://sites.google.com/view/rlbiggym).
 
-Supported by:
+## Citing
+
+If you use BIG-Gym in a publication, please cite referenced libraries:
+
+```
+@article{gu2021braxlines,
+  title={Braxlines: Fast and Interactive Toolkit for RL-driven Behavior Engineering beyond Reward Maximization},
+  author={Gu, Shixiang Shane and Diaz, Manfred and Freeman, Daniel C and Furuta, Hiroki and Ghasemipour, Seyed Kamyar Seyed and Raichuk, Anton and David, Byron and Frey, Erik and Coumans, Erwin and Bachem, Olivier},
+  journal={arXiv preprint arXiv:2110.04686},
+  year={2021}
+}
+@software{brax2021github,
+  author = {C. Daniel Freeman and Erik Frey and Anton Raichuk and Sertan Girgin and Igor Mordatch and Olivier Bachem},
+  title = {Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation},
+  url = {http://github.com/google/brax},
+  version = {0.0.5},
+  year = {2021},
+}
+```
+
+## Organizers
+* [Shixiang Shane Gu](https://sites.google.com/view/gugurus/home) (Google Brain), [Hiroki Furuta](https://frt03.github.io/) (University of Tokyo), [Manfred Diaz](https://manfreddiaz.github.io/) (University of Montreal)
 * [Brax](https://github.com/google/brax)/[Braxlines](https://arxiv.org/abs/2110.04686) teams
 * [NeurIPS 2021 EcoRL workshop](https://sites.google.com/view/ecorl2021/home) organizers

brax/experimental/biggym/__init__.py

@@ -14,13 +14,17 @@
 
 """BIG-Gym: crowd-sourced environments and behaviors."""
 # pylint:disable=protected-access
+# pylint:disable=g-complex-comprehension
+import difflib
 import functools
 import importlib
+import inspect
 from typing import Any, Union, Dict, Callable, Optional
 from brax import envs as brax_envs
 from brax.envs import Env
 from brax.envs import wrappers
-from brax.experimental.biggym.tasks import TASKS
+from brax.experimental.biggym import registry
+from brax.experimental.biggym import tasks
 from brax.experimental.braxlines.envs import obs_indices
 from brax.experimental.composer import components as composer_components
 from brax.experimental.composer import composer
@@ -31,10 +35,18 @@
 
 ROOT_PATH = 'brax.experimental.biggym.registry'
 ENVS = {}
+REGISTRIES = {}
+OPEN_ENDED_TRACKS = ('rl', 'mimax')
+GOAL_ORIENTED_TRACKS = sorted(tasks.TASKS)
+ENVS_BY_TRACKS = dict(
+    open_ended={k: () for k in OPEN_ENDED_TRACKS},
+    goal_oriented={k: () for k in GOAL_ORIENTED_TRACKS},
+)
 
 
 def inspect_env(env_name: str):
   """Inspect env_params of an env (ComposerEnv only)."""
+  assert_exists(env_name)
   if composer_envs.exists(env_name):
     return composer_envs.inspect_env(env_name)
   else:
@@ -45,6 +57,7 @@ def assert_env_params(env_name: str,
                       env_params: Dict[str, Any],
                       ignore_kwargs: bool = True):
   """Inspect env_params of an env (ComposerEnv only)."""
+  assert_exists(env_name)
   if composer_envs.exists(env_name):
     composer_envs.assert_env_params(env_name, env_params, ignore_kwargs)
   else:
@@ -61,12 +74,65 @@ def exists(env_name: str):
   return env_name in list_env()
 
 
-def register(registry_name: str, assert_override: bool = True):
+def assert_exists(env_name: str):
+  """Assert if an environment is registered."""
+  exists_ = exists(env_name)
+  if not exists_:
+    closest = difflib.get_close_matches(env_name, list_env(), n=3)
+    assert 0, f'{env_name} not found. Closest={closest}'
+
+
+def get_func_kwargs(func):
+  """Get keyword args of a function."""
+  # first, unwrap functools.partial. only extra keyword arguments.
+  partial_supported_params = {}
+  while isinstance(func, functools.partial):
+    partial_supported_params.update(func.keywords)
+    func = func.func
+  # secondly, inspect the original function for keyword arguments.
+  fn_params = inspect.signature(func).parameters
+  support_kwargs = any(
+      v.kind == inspect.Parameter.VAR_KEYWORD for v in fn_params.values())
+  supported_params = {
+      k: v.default
+      for k, v in fn_params.items()
+      if v.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD and
+      v.default != inspect._empty
+  }
+  supported_params.update(partial_supported_params)
+  return supported_params, support_kwargs
+
+
+def register_all(verbose: bool = False, **kwargs):
+  """Register all registries."""
+  for registry_name in registry.REGISTRIES:
+    env_names, comp_names, task_env_names = register(registry_name, **kwargs)
+    if verbose:
+      print((f'Registered {registry_name}: '
+             f'{len(env_names)} envs, '
+             f'{len(comp_names)} comps, '
+             f'{len(task_env_names)} task_envs, '))
+
+
+def register(registry_name: str,
+             assert_override: bool = True,
+             optional: bool = True):
   """Register all envs and components."""
-  global ENVS
+  global ENVS, REGISTRIES, ENVS_BY_TRACKS
+
+  assert (optional or registry_name not in REGISTRIES
+         ), f'non-optional register() conflicts: {registry_name}'
+  if registry_name in REGISTRIES:
+    return REGISTRIES[registry_name]
+
   lib = importlib.import_module(f'{ROOT_PATH}.{registry_name}')
   envs = lib.ENVS or {}
   components = lib.COMPONENTS or {}
+  envs = {registry.get_env_name(registry_name, k): v for k, v in envs.items()}
+  components = {
+      registry.get_comp_name(registry_name, k): v
+      for k, v in components.items()
+  }
   task_envs = []
 
   # register environments
@@ -86,7 +152,11 @@ def register(registry_name: str, assert_override: bool = True):
     else:
       # register a standard Env
       ENVS[env_name] = env_module
-    if 'mimax' in env_info.get('tracks', []):
+    tracks = env_info.get('tracks', ['rl'])
+    for track in tracks:
+      assert track in OPEN_ENDED_TRACKS, f'{track} not in {OPEN_ENDED_TRACKS}'
+      ENVS_BY_TRACKS['open_ended'][track] += (env_name,)
+    if 'mimax' in tracks:
       # (MI-Max only) register obs_indices
       for indices_type, indices in env_info.get('obs_indices', {}).items():
         obs_indices.register_indices(env_name, indices_type, indices)
@@ -98,25 +168,30 @@ def register(registry_name: str, assert_override: bool = True):
           comp_name
       ), f'{composer_components.list_components()} contains {comp_name}'
     comp_module = comp_info['module']
-    composer_components.register_component(
+    comp_lib = composer_components.register_component(
         comp_name,
         load_path=f'{ROOT_PATH}.{registry_name}.components.{comp_module}',
         override=True)
+    component_params = get_func_kwargs(comp_lib.get_specs)[0]
     for track in comp_info.get('tracks', []):
-      assert track in TASKS, f'{track} not in {sorted(TASKS)}'
-      track_env_name = f'{track}_{registry_name}_{comp_name}'
-      track_env_module = TASKS[track](
-          component=comp_module,
-          component_params=comp_info.get('component_params', {}))
+      assert (track
+              in GOAL_ORIENTED_TRACKS), f'{track} not in {GOAL_ORIENTED_TRACKS}'
+      track_env_name = tasks.get_task_env_name(track, comp_name)
       if assert_override:
         assert not exists(
             track_env_name), f'{list_env()} contains {track_env_name}'
+      track_env_module = functools.partial(tasks.TASKS[track], comp_name,
+                                           **component_params)
       # register a ComposerEnv
       composer_envs.register_env(
           track_env_name, track_env_module, override=True)
       task_envs += [track_env_name]
+      ENVS_BY_TRACKS['goal_oriented'][track] += (track_env_name,)
 
-  return sorted(envs), sorted(components), sorted(task_envs)
+  assert envs or task_envs, 'no envs registered'
+  REGISTRIES[registry_name] = (sorted(envs), sorted(components),
+                               sorted(task_envs))
+  return REGISTRIES[registry_name]
 
 
 def create(env_name: str = None,
@@ -126,6 +201,7 @@ def create(env_name: str = None,
            batch_size: Optional[int] = None,
            **kwargs) -> Env:
   """Creates an Env with a specified brax system."""
+  assert_exists(env_name)
   if env_name in ENVS:
     env = ENVS[env_name](**kwargs)
     if episode_length is not None:

brax/experimental/biggym/registry/__init__.py

@@ -0,0 +1,29 @@
+# Copyright 2021 The Brax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Registry of BIG-Gym environments and components."""
+
+# keep alphabetical ordering
+REGISTRIES = [
+    'jump',
+    'proant',
+]
+
+
+def get_comp_name(registry_name: str, comp_name: str):
+  return f'{registry_name}__{comp_name}'
+
+
+def get_env_name(registry_name: str, env_name: str):
+  return f'{registry_name}__{env_name}'

brax/experimental/biggym/registry/jump_cheetah/__init__.py → brax/experimental/biggym/registry/jump/__init__.py

@@ -15,8 +15,8 @@
 """Example: an existing Env + a new reward."""
 
 ENVS = dict(
-    jump_cheetah=dict(
-        module='jump_cheetah:JumpCheetah',
+    cheetah=dict(
+        module='cheetah:JumpCheetah',
         tracks=('rl',),
     ),)
 

brax/experimental/biggym/registry/procedural_ant/__init__.py → brax/experimental/biggym/registry/proant/__init__.py

@@ -15,13 +15,13 @@
 """Example: a Component + env rewards."""
 
 ENVS = dict(
-    ant_run_bg=dict(
+    run=dict(
         module='ant:Run',
         tracks=('rl',),
     ),)
 
 COMPONENTS = dict(
-    ant_bg=dict(
+    ant=dict(
         module='ant',
-        tracks=('run',),
+        tracks=('race',),
     ),)

brax/experimental/biggym/registry/proant/components/ant.py

@@ -0,0 +1,17 @@
+# Copyright 2021 The Brax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Procedural ant."""
+# pylint:disable=unused-import
+from brax.experimental.composer.components.pro_ant import get_specs

brax/experimental/biggym/registry/procedural_ant/envs/ant.py → brax/experimental/biggym/registry/proant/envs/ant.py

@@ -13,13 +13,14 @@
 # limitations under the License.
 
 """Ant tasks."""
+from brax.experimental.biggym import registry
 
 
 def Run(num_legs: int = 4):
   return dict(
       components=dict(
           agent1=dict(
-              component='ant_bg',
+              component=registry.get_comp_name('proant', 'ant'),
               component_params=dict(num_legs=num_legs),
               pos=(0, 0, 0),
               reward_fns=dict(