I'm using QDax 0.1.0 on Windows with Jupyter with cpu-only jaxlib. I'm looking at and modifying the map elites notebook. With no modifications, each iteration in the main for-loop takes about 7-8 seconds (looking at mapelites-logs.csv). If I use a custom environment that just does basic jnp operations and returns done after one step, the iteration time only comes down to around 4 seconds. Why can't I get it to something much much faster? I feel like something is being re-jitted. It looks conspicuous in a Google Colab too while it's running. That call-stack preview thing gets long in the bottom of the screen.

Can you explain which parameters are supposed to affect the speed of each iteration? How does num_centroids affect computational cost? How does the size of the action space affect computational cost? How should one pick a batch size?

These were my modification to the map elites notebook.

# change the policy layers
policy_hidden_layer_sizes = (4, 4)
# and re-initialize the policy (not shown)

from brax.envs.env import State
from qdax.environments import QDEnv
from typing import List, Tuple
class MyEnv(QDEnv):

    @property
    def state_descriptor_length(self) -> int:
        raise ValueError("foo")

    @property
    def state_descriptor_name(self) -> str:
        raise ValueError("foo")

    @property
    def state_descriptor_limits(self) -> Tuple[List[float], List[float]]:
        raise ValueError("foo")

    @property
    def behavior_descriptor_length(self) -> int:
        return 3

    @property
    def behavior_descriptor_limits(self) -> Tuple[List[float], List[float]]:
        a_min = [-1. for _ in range(self.behavior_descriptor_length)]
        a_max = [1. for _ in range(self.behavior_descriptor_length)]
        return a_min, a_max

    @property
    def name(self) -> str:
        return "MyEnvFoo"

    @property
    def observation_size(self):
        return 10

    @property
    def action_size(self) -> int:
        return 3

    def reset(self, rng: jnp.ndarray) -> State:
        """Resets the environment to an initial state."""

        obs_init =  jnp.ones((10,))
        
        reward, done = jnp.zeros(2)
        metrics: Dict = {}
        info_init = {"state_descriptor": obs_init}
        return State(None, obs_init, reward, done, metrics, info_init)

    def step(self, state: State, actions) -> State:
        """Run one timestep of the environment's dynamics."""
        
        reward = 1e-6
        done = jnp.array(1.0)
        new_obs = state.obs
        return state.replace(obs=new_obs, reward=reward, done=done)

env = MyEnv(config=None)
# don't use the brax environment
# env = environments.create(env_name, episode_length=episode_length)

Inside play_step_fn, set truncations to None.

Redefine bd_extraction_fn:

def bd_extraction_fn(data, mask):
    # print('actions:', data.actions)
    return data.actions[:,0,:]

The iteration time is still 4 seconds. Thanks for your help. I would love to see this run blazing fast.

Related to issue #74

• all jax tree-based functions are now imported from the tree_util module
• updates the dependencies in requirements.txt and setup.py
• Update the version of Brax in the requirements and in the setup.py

This PR also adds the wrapper CompletedEvalWrapper that used to be in Brax, but has been removed in the most recent versions.

In summary, this wrapper used to be present in Brax, and we rely on some elements of it in 4 algorithms (I think those are: SAC, DIAYN, DADS and TD3), and their tests where all failing for that reason.

So what I did is copying the old wrapper of Brax here, and now all the tests pass.

This is a provisional solution to make it work in the same way as before. Maybe we can improve and adapt our own code to the new structure of QDax, but I think that is for another time/PR.

Hello, I have been working with QDax and I think I encountered an issue regarding which solutions are stored in the archive based on their evaluation scores. That is, it seems like there is a chance that not the highest performant solution will be stored in the archive. https://github.com/adaptive-intelligent-robotics/QDax/blob/ad3953dc904f6dc19b5b4c7f90fa2f8c134128e4/qdax/qd_utils/grid_archive.py#L75

More precisely, please execute the following code and trace which solution is stored in the archive at the end as generated in the line https://github.com/adaptive-intelligent-robotics/QDax/blob/ad3953dc904f6dc19b5b4c7f90fa2f8c134128e4/qdax/qd_utils/grid_archive.py#L78.

The code:

import jax
from qdax.qd_utils import grid_archive
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
params_size =  5
batch_size = 3
grid_shape = (30,30)
min_bd = 0
max_bd = 1
repertoire = grid_archive.Repertoire.create(jax.random.normal(key,shape=(params_size,)), min=min_bd, max=max_bd, grid_shape=grid_shape)

params = jnp.array([[-0.4518846,  -2.0728214,   0.02437184,  0.56900173, -2.0105903 ],
 [ 0.31103376, -0.29348192, -0.27793083, -1.2343968,   1.6130152 ],
 [ 1.997377,   -0.9525061,  -0.57822144,  0.8413021,  -2.02012   ]])

bds = jnp.array([[-0.16706778, -0.5440059 ],
 [-0.47653008, -1.6869655 ],
 [-0.9096347,  -0.07636569]])
objs = jnp.array([2, 10, 2 ]) 

dead = jnp.zeros(batch_size)

repertoire = repertoire.add_to_archive(repertoire = repertoire,
                                 pop_p = params,
                                 bds = bds,
                                 eval_scores = objs,
                                 dead = dead)

The result I get:

bds= [[-0.16706778 -0.5440059 ]
 [-0.47653008 -1.6869655 ]
 [-0.9096347  -0.07636569]]

pop_p= [[-0.4518846  -2.0728214   0.02437184  0.56900173 -2.0105903 ]
 [ 0.31103376 -0.29348192 -0.27793083 -1.2343968   1.6130152 ]
 [ 1.997377   -0.9525061  -0.57822144  0.8413021  -2.02012   ]]

eval_scores= [ 2 10  2]

bd_insertion= [0 0 0]

repertoire.archive= [[ 1.997377   -0.9525061  -0.57822144  0.8413021  -2.02012   ]
 [ 0.          0.          0.          0.          0.        ]
 [ 0.          0.          0.          0.          0.        ]
 .
 .
 .
 [ 0.          0.          0.          0.          0.        ]
 [ 0.          0.          0.          0.          0.        ]
 [ 0.          0.          0.          0.          0.        ]
 [ 0.          0.          0.          0.          0.        ]]

You can change the behavioural descriptors to any other values that would map to the same position and the same issue still exists. That is, this doesn't seem to happen on position 0 of the archive array only.

This PR is related to #82

It introduces:

• the introduction of the QualityPGEmitter to decompose the PGAMEEmitter
• refactoring of PGAMEEmitter. Now, PGAMEEmitter = MultiEmitter(QualityPGEmitter, MixingEmitter)
• the algorithm QDPG, through the use of the QDPGEmitter, with QDPGEmitter = MultiEmitter(DiversityPGEmitter, QualityPGEmitter, MixingEmitter)
• test for QDPG
• notebook for QDPG
• updated documentation
• updated MultiEmitter to give the option to use all data when updating the emitters (important for the PGEmitters as they both need the transitions for all policies)

Potential discussion:

• there is a small difference in the way PGAMEEmitter and QDPGEmitter are defined, in particular the config. Reviewers should have a look at both and tell me what they prefer. Also unification is better, they are also reasons to have different style (defining the qdpg config in the same way as pgame would make it very messy; but having pgame config like qdpg config would make it a bit less straitforward and introduce a small change in the api).
• isn't it time to get rid of the MixingEmitter? My personal answer is yes but I just felt we should keep it for another PR to keep everything clear (also, I did not wanted to change PGAME too much in a single PR). If everyone agrees, I'll create an issue with this task.

Remaining tasks:

• [x] double check PGAME perfs
• [x] get some QDPG perfs
• [x] wait for MultiEmitter PR to be merged in develop first

Delayed update of the actor policy when updating the models (actor and critic) in PGA

Before update:

After update:

Performance is tested on walker_uni task. Performance seems to improve (albeit only testing on one seed). Coverage relatively unaffected hence improvement mainly in quality, which makes sense given that this is a fix for the actor policy update.

switched depreciated dependency sklearn -> scikit-learn

Related issues: #130

The PR replaces the depreciated sklearn package with scikit-learn.

Checks

• [x] a clear description of the PR has been added
• [ ] sufficient tests have been written
• [ ] relevant section added to the documentation
• [ ] example notebook added to the repo
• [ ] clean docstrings and comments have been written
• [ ] if any issue/observation has been discovered, a new issue has been opened

Future improvements

N/A

When installing qdax, sklearn package raises an exception due to the depreciation of sklearn in favor of scikit-learn (see details on pypi).

(.venv) [email protected] ~/G/yuyu> python3 -m pip install -U --no-cache-dir qdax
Collecting qdax
  Downloading qdax-0.2.0-py3-none-any.whl (155 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 155.3/155.3 kB 16.1 MB/s eta 0:00:00
Requirement already satisfied: jax>=0.3.16 in ./.venv/lib/python3.10/site-packages (from qdax) (0.3.25)
Requirement already satisfied: numpy>=1.22.3 in ./.venv/lib/python3.10/site-packages (from qdax) (1.23.5)
Requirement already satisfied: scipy>=1.8.0 in ./.venv/lib/python3.10/site-packages (from qdax) (1.9.3)
Requirement already satisfied: jaxlib>=0.3.15 in ./.venv/lib/python3.10/site-packages (from qdax) (0.3.25)
Collecting gym>=0.23.1
  Downloading gym-0.26.2.tar.gz (721 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 721.7/721.7 kB 15.9 MB/s eta 0:00:00
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Preparing metadata (pyproject.toml) ... done
Collecting scikit-learn>=1.0.2
  Downloading scikit_learn-1.1.3-cp310-cp310-macosx_10_9_x86_64.whl (8.7 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 8.7/8.7 MB 21.3 MB/s eta 0:00:00
Collecting brax>=0.0.15
  Downloading brax-0.0.15-py3-none-any.whl (372 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 372.3/372.3 kB 31.0 MB/s eta 0:00:00
Collecting flax<0.6.2,>=0.6
  Downloading flax-0.6.1-py3-none-any.whl (185 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 185.6/185.6 kB 126.9 MB/s eta 0:00:00
Collecting sklearn
  Downloading sklearn-0.0.post1.tar.gz (3.6 kB)
  Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error

  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [18 lines of output]
      The 'sklearn' PyPI package is deprecated, use 'scikit-learn'
      rather than 'sklearn' for pip commands.

      Here is how to fix this error in the main use cases:
      - use 'pip install scikit-learn' rather than 'pip install sklearn'
      - replace 'sklearn' by 'scikit-learn' in your pip requirements files
        (requirements.txt, setup.py, setup.cfg, Pipfile, etc ...)
      - if the 'sklearn' package is used by one of your dependencies,
        it would be great if you take some time to track which package uses
        'sklearn' instead of 'scikit-learn' and report it to their issue tracker
      - as a last resort, set the environment variable
        SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=True to avoid this error

      More information is available at
      https://github.com/scikit-learn/sklearn-pypi-package

      If the previous advice does not cover your use case, feel free to report it at
      https://github.com/scikit-learn/sklearn-pypi-package/issues/new
      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.

The wrapper alters the reset function to return the default initial joint positions and zero velocity on the joints. Some issues remain due to the underlying brax envs which do not have uniform api for _get_obs.

• tests have been implemented, fails the humanoid, walker and hopper tests at the moment because of the issue mentioned above.

Add a base Multi-Emitter implementation of a batch of Emitters. No choice strategy is used, all sub emitters are called and the proportion are defined in the script when defining the sub emitters. We will consider adding "strategic" layers later, to manage the choice of emitters or the proportion. This is left for a future PR.

Exceptionally, we'll leave tests and documentation for another PR.

TODOs:

• [x] Tests
• [x] Documentation

Hi :)

The current PGAME Replay Buffer is using jax.lax.dynamic_update_slice to add new transition to the replay buffer. However, this is not acting like a circular buffer, meaning that if a batch contain more transitions than the size remaining in the buffer, it would delete the more recent transitions instead of the oldest ones.

Hi :)

It seems that the current implementation of PGAME is adding to the Replay-Buffer all the transitions that are collected in the environment, meaning also transitions that occur after the individual is dead and the environment returned done = 1. I only run initial tests but this seems to slightly impact the PGAME algorithm performance.

Hello,

Let me start by saying that I am a fan of your work here. I have recently open-sourced by GNN-based meta-learning method for optimization. I have applied it to the sparse index-tracking problem from real-world (after an initial benchmarking on Schwefel function), and it seems to outperform Fast CMA-ES significantly both in terms of producing robust solutions on the blind test set and also in terms of time (total duration and iterations) and space complexity. I include the link to my repository here, in case you would consider adding the method or the benchmarking problem to your repository. Note: GNN, which learns how to generate populations of solutions at each iteration, is trained using gradients retrieved from the loss function, as opposed to black-box ones.

Sincerely, K

This PR introduces:

• neigbors -> neighbors

Checks

• [x] a clear description of the PR has been added
• [ ] sufficient tests have been written
• [ ] relevant section added to the documentation
• [ ] example notebook added to the repo
• [ ] clean docstrings and comments have been written
• [ ] if any issue/observation has been discovered, a new issue has been opened

Future improvements

[List here potential observations made and/or improvements that could be made in the future. If relevant, open issues for those.]

We can do the automatically now with the updated brax version which takes this into account. Also something to add to the caveats section, that the reset of the trap wrapper (any wrapper that adds new bodies) changes the observations size/dimensions.

It is incoherent the way fitnesses and descriptors are ordered in repertoire addition compared to the rest of the codebase. This is very minor but we should still consider fixing this in the future.

Following up issue #65.

Multi-devices support has been resolved and can be used with TPU; but not with docker at the moment. Should we add TPU support for docker containers?