Refactors acme.jax.networks.rescaling to convert hk modules to pure functions.
Resolves
"""Rescaling layers (e.g. to match action specs)."""
+import dataclasses
+
from acme import specs
-import haiku as hk
from jax import lax
import jax.numpy as jnp
-class ClipToSpec(hk.Module):
- """Sonnet module clipping inputs to within a BoundedArraySpec."""
-
- def __init__(self, spec: specs.BoundedArray, name: str = 'clip_to_spec'):
- super().__init__(name=name)
- self._min = spec.minimum
- self._max = spec.maximum
+@dataclasses.dataclass
+class ClipToSpec:
+ """Clips inputs to within a BoundedArraySpec."""
+ spec: specs.BoundedArray
def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
- return jnp.clip(inputs, self._min, self._max)
-
+ return jnp.clip(inputs, self.spec.minimum, self.spec.maximum)
-class RescaleToSpec(hk.Module):
- """Sonnet module rescaling inputs in [-1, 1] to match a BoundedArraySpec."""
- def __init__(self, spec: specs.BoundedArray, name: str = 'rescale_to_spec'):
- super().__init__(name=name)
- self._scale = spec.maximum - spec.minimum
- self._offset = spec.minimum
+@dataclasses.dataclass
+class RescaleToSpec:
+ """Rescales inputs in [-1, 1] to match a BoundedArraySpec."""
+ spec: specs.BoundedArray
def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
+ scale = self.spec.maximum - self.spec.minimum
+ offset = self.spec.minimum
inputs = 0.5 * (inputs + 1.0) # [0, 1]
- output = inputs * self._scale + self._offset # [minimum, maximum]
-
+ output = inputs * scale + offset # [minimum, maximum]
return output
-class TanhToSpec(hk.Module):
- """Sonnet module squashing real-valued inputs to match a BoundedArraySpec."""
-
- def __init__(self, spec: specs.BoundedArray, name: str = 'tanh_to_spec'):
- super().__init__(name=name)
- self._scale = spec.maximum - spec.minimum
- self._offset = spec.minimum
+@dataclasses.dataclass
+class TanhToSpec:
+ """Squashes real-valued inputs to match a BoundedArraySpec."""
+ spec: specs.BoundedArray
def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
+ scale = self.spec.maximum - self.spec.minimum
+ offset = self.spec.minimum
inputs = lax.tanh(inputs) # [-1, 1]
inputs = 0.5 * (inputs + 1.0) # [0, 1]
- output = inputs * self._scale + self._offset # [minimum, maximum]
+ output = inputs * scale + offset # [minimum, maximum]
return output
