Fix enum member py version issue.

PiperOrigin-RevId: 829322405
Change-Id: I471abb204907e0337a53cd74877dd9e7d1ccc446

Author: btaba

brax/training/agents/es/train.py

@@ -19,6 +19,7 @@
 
 import enum
 import functools
+import sys
 import time
 from typing import Any, Callable, Dict, Optional, Tuple
 
@@ -40,6 +41,11 @@
 Metrics = types.Metrics
 InferenceParams = Tuple[running_statistics.NestedMeanStd, Params]
 
+if sys.version_info >= (3, 11):
+  from enum import member as enum_member
+else:
+  enum_member = lambda x: x
+
 
 @flax.struct.dataclass
 class TrainingState:
@@ -67,9 +73,9 @@ def wierstra(x: jnp.ndarray) -> jnp.ndarray:
 
 
 class FitnessShaping(enum.Enum):
-  ORIGINAL = enum.member(functools.partial(lambda x: x))
-  CENTERED_RANK = enum.member(functools.partial(centered_rank))
-  WIERSTRA = enum.member(functools.partial(wierstra))
+  ORIGINAL = enum_member(functools.partial(lambda x: x))
+  CENTERED_RANK = enum_member(functools.partial(centered_rank))
+  WIERSTRA = enum_member(functools.partial(wierstra))
 
 
 # TODO(eorsini): Pass the network as argument.