deps: replace treevalue with optree

[Benjamin-eecs]

Description

• close [BUG] Move away from dm-tree broke compatibility with Jax #246
• further speedup
• support custom node type registration

Motivation and Context

Initial Test Results:

Test gym

TreeValue:

OpTree:

Test dmc

TreeValue:

OpTree:

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Put an x in all the boxes that apply:

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Example (update in the folder of example)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.
• I have reformatted the code using make format (required)
• I have checked the code using make lint (required)
• I have ensured make bazel-test pass. (required)

cc @XuehaiPan

[HansBug]

We are really really glad to see @XuehaiPan @Benjamin-eecs bring some more advise to treevalue. This will make treevalue better and better, because we can add all the features you come up with to treevalue in a very short time, which is faster than you can imagine. We hope to have more in-depth exchanges with the optree team in order to significantly improve the experience of using treevalue. 😄 Love U ❤️

In order to express our gratitude, we suggest that optree should take the EASE OF USE seriously, such as whether it is possible to attach additional methods to the tree data structure, and whether it can be programmed in OOP without great overhead. Treevalue can easily acheive this, and its performance is not low at all. We think this is very important for an AI developer-oriented package, because users will use actual usage to show which one is better. Ease of use can have a profound impact on a long-term maintained library.

[XuehaiPan]

we suggest that optree should take the EASE OF USE seriously, such as whether it is possible to attach additional methods to the tree data structure, and whether it can be programmed in OOP without great overhead.

@HansBug Thanks for your advice. We may consider this in a future release of optree. Currently, we think that always returning Python native types is the top priority of the optree package support.

Currently optree only supports primitive tree operations, and native types in and native types out.

In [1]: from collections import *

In [2]: import optree
   ...: from optree.typing import PyTree, T

In [3]: tree = {'a': 1, 'b': {'c': [2, (3, 4)], 'd': deque([4, 5])}, 'e': OrderedDict([('f', 6), ('g', 7)])}

In [4]: def tree_sum(tree: PyTree[T]) -> T:
   ...:     return optree.tree_reduce(lambda x, y: x + y, tree, 0)
   ...:

In [5]: tree_sum(tree)
Out[5]: 32

In [6]: def tree_add(tree: PyTree[T], other: PyTree[T]) -> PyTree[T]:
   ...:     return optree.tree_map(lambda x, y: x + y, tree, other)
   ...:

In [7]: other = optree.tree_map(lambda x: 2 * x, tree)

In [8]: other
Out[8]: {'a': 2, 'b': {'c': [4, (6, 8)], 'd': deque([8, 10])}, 'e': OrderedDict([('f', 12), ('g', 14)])}

In [9]: tree_add(tree, other)
Out[9]: {'a': 3, 'b': {'c': [6, (9, 12)], 'd': deque([12, 15])}, 'e': OrderedDict([('f', 18), ('g', 21)])}

optree can also be extended by creating OOP wrappers implemented with dunder methods (e.g., __pos__, __add__, __mul__) using tree_map. It still needs some effort to do, such as tree broadcast when the tree structures mismatch. In this scope, I think the usage of treevalue is more intuitive. FYI, you could check out pytorch-labs/tensordict, another dict-based container with similar functionality.