unienv 0.0.1b1__py3-none-any.whl → 0.0.1b3__py3-none-any.whl

unienv_interface/space/space_utils/batch_utils.py

@@ -5,6 +5,7 @@ import copy
 from functools import singledispatch
 from typing import Optional, Any, Iterable, Iterator
 from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType, ArrayAPIGetIndex, ArrayAPISetIndex
+from unienv_interface.backends.numpy import NumpyComputeBackend
 import numpy as np
 
 from ..spaces import *
@@ -178,6 +179,19 @@ def _reshape_batch_size_tuple(
         device=space.device,
     )
 
+@reshape_batch_size.register(BatchedSpace)
+def _reshape_batch_size_batched(
+    space: BatchedSpace,
+    old_batch_shape: typing.Tuple[int],
+    new_batch_shape: typing.Tuple[int]
+) -> BatchedSpace:
+    assert len(old_batch_shape) <= len(space.batch_shape) and old_batch_shape == space.batch_shape[:len(old_batch_shape)], \
+        f"Expected the old batch shape to be a prefix of the current batch shape, but got old {old_batch_shape} != current {space.batch_shape}"
+    return BatchedSpace(
+        single_space=space.single_space,
+        batch_shape=new_batch_shape + space.batch_shape[len(old_batch_shape):]
+    )
+
 def reshape_batch_size_in_data(
     backend : ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType],
     data : Any, 
@@ -194,6 +208,10 @@ def reshape_batch_size_in_data(
         assert data.shape[:len(old_batch_shape)] == old_batch_shape, \
             f"Expected the beginning of the shape to match the old batch shape, but got {data.shape[:len(old_batch_shape)]} != {old_batch_shape}"
         data = data.reshape(new_batch_shape + data.shape[len(old_batch_shape):])
+    elif isinstance(data, np.ndarray) and data.dtype == object:
+        assert data.shape[:len(old_batch_shape)] == old_batch_shape, \
+            f"Expected the beginning of the shape to match the old batch shape, but got {data.shape[:len(old_batch_shape)]} != {old_batch_shape}"
+        data = data.reshape(new_batch_shape + data.shape[len(old_batch_shape):])
     elif isinstance(data, GraphInstance):
         assert data.n_nodes.shape[:len(old_batch_shape)] == old_batch_shape, \
             f"Expected the beginning of the n_nodes shape to match the old batch shape, but got {data.n_nodes.shape[:len(old_batch_shape)]} != {old_batch_shape}"
@@ -301,6 +319,12 @@ def _swap_batch_dims_tuple(space: TupleSpace, dim1: int, dim2: int):
         device=space.device,
     )
 
+@swap_batch_dims.register(BatchedSpace)
+def _swap_batch_dims_batched(space: BatchedSpace, dim1: int, dim2: int):
+    return BatchedSpace(
+        single_space=space.single_space,
+        batch_shape=_shape_transpose(space.batch_shape, dim1, dim2)
+    )
 
 def swap_batch_dims_in_data(
     backend : ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType], 
@@ -310,6 +334,8 @@ def swap_batch_dims_in_data(
 ) -> Any:
     if backend.is_backendarray(data):
         return _tensor_transpose(backend, data, dim1, dim2)
+    elif isinstance(data, np.ndarray) and data.dtype != object:
+        return _tensor_transpose(NumpyComputeBackend, data, dim1, dim2)
     elif isinstance(data, GraphInstance):
         return GraphInstance(
             n_nodes=_tensor_transpose(backend, data.n_nodes, dim1, dim2),
@@ -367,6 +393,10 @@ def _batch_size_tuple(space: TupleSpace):
 
     return len(space.spaces)
 
+@batch_size.register(BatchedSpace)
+def _batch_size_batched(space: BatchedSpace):
+    return space.batch_shape[0] if len(space.batch_shape) > 0 else None
+
 def batch_size_data(data: Any) -> Optional[int]:
     if hasattr(data, "shape"):
         return data.shape[0] if len(data.shape) > 0 else None
@@ -451,6 +481,21 @@ def _batch_space_tuple(space: TupleSpace, n: int = 1):
         device=space.device,
     )
 
+@batch_space.register(BatchedSpace)
+def _batch_space_batched(space: BatchedSpace, n: int = 1):
+    return BatchedSpace(
+        single_space=space.single_space,
+        batch_shape=(n,) + space.batch_shape,
+    )
+
+@batch_space.register(TextSpace)
+@batch_space.register(UnionSpace)
+def _batch_space_text(space: typing.Union[TextSpace, UnionSpace], n: int = 1):
+    return BatchedSpace(
+        space,
+        batch_shape=(n,),
+    )
+
 @singledispatch
 def batch_differing_spaces(spaces: typing.Sequence[Space], device : Optional[Any] = None) -> Space:
     assert len(spaces) > 0, "Expects a non-empty list of spaces"
@@ -642,6 +687,16 @@ def _unbatch_spaces_tuple(space: TupleSpace):
             device=space.device,
         )
 
+@unbatch_spaces.register(BatchedSpace)
+def _unbatch_spaces_batched(space: BatchedSpace):
+    assert len(space.batch_shape) > 0, "Expected BatchedSpace to be batched, but it is not."
+    unbatched_space = space.single_space if len(space.batch_shape) == 1 else BatchedSpace(
+        single_space=space.single_space,
+        batch_shape=space.batch_shape[1:],
+    )
+    for i in range(space.batch_shape[0]):
+        yield unbatched_space
+
 def iterate(space: Space, items: Any) -> Iterator:
     for i in range(batch_size_data(items)):
         yield get_at(space, items, i)
@@ -677,6 +732,10 @@ def _get_at_dict(space: DictSpace, items: typing.Mapping[str, Any], index : Arra
 def _get_at_tuple(space: TupleSpace, items: typing.Tuple[Any, ...], index : ArrayAPIGetIndex):    
     return tuple(get_at(subspace, item, index) for (subspace, item) in zip(space.spaces, items))
 
+@get_at.register(BatchedSpace)
+def _get_at_batched(space: BatchedSpace, items: np.ndarray, index: ArrayAPIGetIndex) -> typing.Union[np.ndarray, Any]:
+    return items[index]
+
 @singledispatch
 def set_at(
     space: Space, items: Any, index: ArrayAPISetIndex, value: Any
@@ -763,6 +822,17 @@ def _set_at_tuple(
         for i, subspace in enumerate(space.spaces)
     )
 
+@set_at.register(BatchedSpace)
+def _set_at_batched(
+    space: BatchedSpace,
+    items: np.ndarray,
+    index: ArrayAPISetIndex,
+    value: typing.Union[np.ndarray, Any],
+) -> np.ndarray:
+    new_data = items.copy()
+    new_data[index] = value
+    return new_data
+
 @singledispatch
 def concatenate(
     space: Space, items: Iterable[Any], axis : int = 0,
@@ -865,3 +935,16 @@ def _concatenate_tuple(
         )
     
     return tuple(items)
+
+@concatenate.register(BatchedSpace)
+def _concatenate_batched(
+    space: BatchedSpace, items: Iterable, axis: int = 0
+) -> Any:
+    items = list(items)
+    if len(items) == 0:
+        return np.array([], dtype=object)
+    if isinstance(items[0], np.ndarray) and items[0].dtype == object:
+        return np.concatenate(items, axis=axis)
+    else:
+        assert axis == 0, "Expected axis to be 0 when concatenating non-numpy arrays"
+        return np.asarray(items, dtype=object)

unienv_interface/space/space_utils/construct_utils.py

@@ -0,0 +1,216 @@
+import typing
+from copy import deepcopy
+from functools import singledispatch
+from typing import Optional, Any, Iterable, Iterator, Sequence, Tuple, Literal, Mapping
+import numpy as np
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+from unienv_interface.space.spaces import *
+
+__all__ = [
+    'construct_space_from_data_stream',
+    'construct_space_from_data',
+]
+
+def construct_space_from_data_stream(
+    data : Iterable[Any],
+    backend : ComputeBackend,
+    add_bounds : bool = True
+) -> Space:
+    """Construct a space from a stream of data.
+
+    Args:
+        data (Iterable[Any]): An iterable stream of data samples.
+        backend (ComputeBackend): The compute backend to use for array operations.
+
+    Returns:
+        Space: The constructed space.
+    """
+    candidate = None
+    for d in data:
+        candidate = construct_space_from_data(d, backend, candidate, add_bounds=add_bounds)
+    return candidate
+
+def construct_space_from_data(
+    data : Any,
+    backend : ComputeBackend,
+    candidate_space : Optional[Space] = None,
+    add_bounds : bool = False
+) -> Space:
+    if backend.is_backendarray(data) and backend.dtype_is_boolean(data.dtype):
+        assert candidate_space is None or (
+            isinstance(candidate_space, BinarySpace) and
+            candidate_space.shape == data.shape and
+            (candidate_space.dtype is None or candidate_space.dtype == data.dtype) and
+            (candidate_space.device is None or candidate_space.device == backend.device(data))
+        )
+
+        return BinarySpace(
+            backend,
+            data.shape,
+            dtype=data.dtype,
+            device=backend.device(data)
+        )
+    elif backend.is_backendarray(data):
+        assert candidate_space is None or (
+            isinstance(candidate_space, (BoxSpace, DynamicBoxSpace)) and
+            candidate_space.dtype == data.dtype and
+            (candidate_space.device is None or candidate_space.device == backend.device(data))
+        )
+        if candidate_space is None:
+            return BoxSpace(
+                backend,
+                low=data if add_bounds else -backend.inf,
+                high=data if add_bounds else backend.inf,
+                shape=data.shape,
+                dtype=data.dtype,
+                device=backend.device(data)
+            )
+        elif isinstance(candidate_space, BoxSpace):
+            if candidate_space.shape != data.shape:
+                assert len(candidate_space.shape) == len(data.shape)
+                new_low_shape, new_high_shape = _dynamic_box_find_shape(
+                    candidate_space.shape,
+                    candidate_space.shape,
+                    data.shape
+                )
+                broadcast_shape = _dynamic_box_get_broadcast_shape(new_low_shape, new_high_shape)
+
+                return DynamicBoxSpace(
+                    backend,
+                    *(_dynamic_box_update_bounds(
+                        backend,
+                        broadcast_shape,
+                        candidate_space._low,
+                        candidate_space._high,
+                        data
+                    ) if add_bounds else (candidate_space._low, candidate_space._high)),
+                    shape_low=new_low_shape,
+                    shape_high=new_high_shape,
+                    dtype=data.dtype,
+                    device=backend.device(data)
+                )
+        else: # elif isinstance(candidate_space, DynamicBoxSpace):
+            new_low_shape, new_high_shape = _dynamic_box_find_shape(
+                candidate_space.shape_low,
+                candidate_space.shape_high,
+                data.shape
+            )
+            broadcast_shape = _dynamic_box_get_broadcast_shape(new_low_shape, new_high_shape)
+
+            return DynamicBoxSpace(
+                backend,
+                *(_dynamic_box_update_bounds(
+                    backend,
+                    broadcast_shape,
+                    candidate_space._low,
+                    candidate_space._high,
+                    data
+                ) if add_bounds else (candidate_space._low, candidate_space._high)),
+                shape_low=new_low_shape,
+                shape_high=new_high_shape,
+                dtype=data.dtype,
+                device=backend.device(data)
+            )
+    elif isinstance(data, Mapping):
+        assert candidate_space is None or (
+            isinstance(candidate_space, DictSpace)
+            and set(candidate_space.spaces.keys()) == set(data.keys())
+        )
+        if candidate_space is None:
+            spaces = {k: construct_space_from_data(v, backend) for k, v in data.items()}
+        else:
+            spaces = {
+                k: construct_space_from_data(v, backend, candidate_space.spaces[k])
+                for k, v in data.items()
+            }
+        return DictSpace(backend, spaces)
+    elif isinstance(data, str):
+        assert candidate_space is None or isinstance(candidate_space, TextSpace)
+        max_length = len(data)
+        if candidate_space is not None:
+            max_length = max(max_length, candidate_space.max_length)
+        if not add_bounds:
+            max_length = max(max_length, 4096) # Arbitrary large length if not adding bounds
+        return TextSpace(
+            backend,
+            max_length=max_length
+        )
+    elif isinstance(data, Sequence):
+        assert candidate_space is None or (
+            isinstance(candidate_space, TupleSpace)
+            and len(candidate_space.spaces) == len(data)
+        )
+        if candidate_space is None:
+            spaces = [construct_space_from_data(d, backend) for d in data]
+        else:
+            spaces = [
+                construct_space_from_data(d, backend, candidate_space.spaces[i])
+                for i, d in enumerate(data)
+            ]
+        return TupleSpace(backend, spaces)
+    else:
+        raise ValueError(f"Unsupported data type for space construction: {type(data)}")
+
+def _dynamic_box_find_shape(
+    shape_low : Sequence[int],
+    shape_high : Sequence[int],
+    data_shape : Sequence[int]
+) -> Tuple[Sequence[int], Sequence[int]]:
+    assert len(shape_low) == len(shape_high) == len(data_shape)
+    new_shape_low = list(shape_low)
+    new_shape_high = list(shape_high)
+    for i in range(len(data_shape)):
+        new_shape_low[i] = min(shape_low[i], data_shape[i])
+        new_shape_high[i] = max(shape_high[i], data_shape[i])
+    return new_shape_low, new_shape_high
+
+def _dynamic_box_get_broadcast_shape(
+    shape_low : Sequence[int],
+    shape_high : Sequence[int],
+) -> Sequence[int]:
+    assert len(shape_low) == len(shape_high)
+    broadcast_shape = []
+    for low, high in zip(shape_low, shape_high):
+        if low == high:
+            broadcast_shape.append(low)
+        else:
+            broadcast_shape.append(1)  # Use 1 to indicate dynamic dimension
+    return tuple(broadcast_shape)
+
+def reshape_to_broadcastable(
+    backend : ComputeBackend,
+    target_shape : Sequence[int],
+    array : BArrayType,
+    method : Literal['min', 'max'] = 'min'
+) -> BArrayType:
+    if array.shape == target_shape:
+        return array
+    else:
+        assert len(array.shape) == len(target_shape)
+        for i, (t_dim, a_dim) in enumerate(zip(target_shape, array.shape)):
+            assert t_dim == a_dim or t_dim == 1 or a_dim == 1
+            if t_dim == a_dim or a_dim == 1:
+                continue
+            else:
+                if method == 'min':
+                    array = backend.min(array, axis=i, keepdims=True)
+                else:
+                    array = backend.max(array, axis=i, keepdims=True)
+        return array
+
+def _dynamic_box_update_bounds(
+    backend : ComputeBackend,
+    target_shape : Sequence[int],
+    current_low : BArrayType,
+    current_high : BArrayType,
+    new_data : BArrayType
+) -> Tuple[BArrayType, BArrayType]:
+    new_low = backend.minimum(
+        reshape_to_broadcastable(backend, target_shape, current_low, method='min'),
+        reshape_to_broadcastable(backend, target_shape, new_data, method='min')
+    )
+    new_high = backend.maximum(
+        reshape_to_broadcastable(backend, target_shape, current_high, method='max'),
+        reshape_to_broadcastable(backend, target_shape, new_data, method='max')
+    )
+    return new_low, new_high

unienv_interface/space/space_utils/serialization_utils.py

@@ -160,7 +160,7 @@ def _text_space_to_json(space: TextSpace) -> typing.Dict[str, Any]:
         "type": "TextSpace",
         "min_length": space.min_length,
         "max_length": space.max_length,
-        "charset": "".join(space.charset)
+        "charset": "".join(space.charset) if space.charset is not None else None,
     }
 
 @json_to_space.register(TextSpace)
@@ -219,3 +219,18 @@ def _json_to_union_space(json_data: typing.Dict[str, Any], map_backend: ComputeB
     spaces = [json_to_space(s, map_backend, map_device) for s in json_data["spaces"]]
     return UnionSpace(map_backend, spaces, device=map_device)
 
+@space_to_json.register(BatchedSpace)
+def _batched_space_to_json(space: BatchedSpace) -> typing.Dict[str, Any]:
+    return {
+        "type": "BatchedSpace",
+        "single_space": space_to_json(space.single_space),
+        "batch_shape": space.batch_shape,
+    }
+
+@json_to_space.register(BatchedSpace)
+def _json_to_batched_space(json_data: typing.Dict[str, Any], map_backend: ComputeBackend, map_device: Optional[BDeviceType]) -> BatchedSpace:
+    single_space = json_to_space(json_data["single_space"], map_backend, map_device)
+    return BatchedSpace(
+        single_space,
+        batch_shape=json_data["batch_shape"]
+    )

unienv_interface/space/spaces/__init__.py

@@ -7,6 +7,7 @@ from .graph import GraphSpace, GraphInstance
 from .text import TextSpace
 from .tuple import TupleSpace
 from .union import UnionSpace
+from .batched import BatchedSpace
 
 __all__ = [
     "Space",
@@ -18,5 +19,6 @@ __all__ = [
     "GraphInstance",
     "TextSpace",
     "TupleSpace",
-    "UnionSpace"
+    "UnionSpace",
+    "BatchedSpace",
 ]

unienv_interface/space/spaces/batched.py

@@ -0,0 +1,90 @@
+"""Implementation of a space consisting of finitely many elements."""
+from typing import Any, Generic, Iterable, SupportsFloat, Mapping, Sequence, TypeVar, Optional, Tuple, Type, Literal, Union, Callable
+import numpy as np
+from ..space import Space, SpaceDataT
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+
+class BatchedSpace(Space[np.ndarray, BDeviceType, BDtypeType, BRNGType]):
+    """
+    This space represents a batch of 
+    """
+    def __init__(
+        self, 
+        single_space : Space[SpaceDataT, BDeviceType, BDtypeType, BRNGType],
+        batch_shape: Sequence[int],
+    ):
+        assert len(batch_shape) > 0, "Batch shape must be a non-empty sequence"
+        batch_shape = tuple(int(dim) for dim in batch_shape)  # This changes any np types to int
+        super().__init__(
+            backend=single_space.backend,
+            shape=batch_shape + ((),) if single_space.shape is None else batch_shape + single_space.shape,
+            device=single_space.device,
+            dtype=single_space.dtype,
+        )
+
+        self.batch_shape = batch_shape
+        self.single_space = single_space
+
+    def to(self, backend = None, device = None):
+        return BatchedSpace(
+            self.single_space.to(backend=backend, device=device),
+            self.batch_shape
+        )
+    
+    def sample(self, rng : BRNGType) -> Tuple[
+        BRNGType, BArrayType
+    ]:
+        flat_shape = np.prod(self.batch_shape)
+        samples = []
+        for i in range(flat_shape):
+            rng, single_sample = self.single_space.sample(rng)
+            samples.append(single_sample)
+        return rng, np.asarray(samples, dtype=object).reshape(self.batch_shape)
+    
+    def create_empty(self):
+        flat_shape = np.prod(self.batch_shape)
+        empties = [self.single_space.create_empty() for _ in range(flat_shape)]
+        return np.asarray(empties, dtype=object).reshape(self.batch_shape)
+    
+    def is_bounded(self, manner = "both"):
+        return self.single_space.is_bounded(manner=manner)
+
+    def contains(self, x: BArrayType) -> bool:
+        def is_contained_func(x):
+            return self.single_space.contains(x)
+        for _dim in reversed(self.batch_shape):
+            def new_is_contained_func(x):
+                if (not isinstance(x, np.ndarray)) or (not x.dtype == object):
+                    return False
+                if len(x) != _dim:
+                    return False
+                return all(is_contained_func(xi) for xi in x)
+            is_contained_func = new_is_contained_func
+        return is_contained_func(x)
+
+    def get_repr(
+        self, 
+        abbreviate = False,
+        include_backend = True, 
+        include_device = True, 
+        include_dtype = True
+    ):
+        ret = f"BatchedSpace({self.single_space}, batch_shape={self.batch_shape}"
+        if include_backend:
+            ret += f", {self.backend}"
+        if include_device:
+            ret += f", {self.device}"
+        if include_dtype:
+            ret += f", {self.dtype}"
+        ret += ")"
+        return ret
+
+    def __eq__(self, other: Any) -> bool:
+        """Check whether `other` is equivalent to this instance."""
+        return isinstance(other, BatchedSpace) and self.backend == other.backend and self.batch_shape == other.batch_shape and self.single_space == other.single_space
+    
+    def data_to(self, data, backend = None, device = None):
+        if isinstance(data, np.ndarray) and data.dtype == object:
+            return tuple(self.data_to(d, backend=backend, device=device) for d in data)
+        else:
+            return self.single_space.data_to(data, backend=backend, device=device)

unienv_interface/space/spaces/binary.py

@@ -13,7 +13,6 @@ class BinarySpace(Space[BArrayType, BDeviceType, BDtypeType, BRNGType]):
         device : Optional[BDeviceType] = None,
     ):
         assert dtype is None or backend.dtype_is_boolean(dtype), f"Invalid dtype {dtype}"
-        assert len(shape) > 0, "Shape must be a non-empty sequence"
         
         assert all(
             np.issubdtype(type(dim), np.integer) for dim in shape

unienv_interface/space/spaces/box.py

@@ -53,31 +53,20 @@ class BoxSpace(Space[BArrayType, BDeviceType, BDtypeType, BRNGType]):
             dttype_iinfo = backend.iinfo(dtype)
             dtype_min = dttype_iinfo.min
             dtype_max = dttype_iinfo.max
-            if isinstance(low, int):
-                if low == backend.inf or low == -backend.inf:
-                    _low = dtype_min if low == -backend.inf else dtype_max
-                else:
-                    _low = low
-                _low = backend.full([1] * len(shape), _low, dtype=dtype, device=device)
-            else:
-                _low = backend.astype(low, dtype)
-            if isinstance(high, int):
-                if high == backend.inf or high == -backend.inf:
-                    _high = dtype_max if high == backend.inf else dtype_min
-                else:
-                    _high = high
-                _high = backend.full([1] * len(shape), _high, dtype=dtype, device=device)
-            else:
-                _high = backend.astype(high, dtype)
+
+            if isinstance(low, (int, float)) and (low == backend.inf or low == -backend.inf):
+                low = dtype_min if low == -backend.inf else dtype_max
+            if isinstance(high, (int, float)) and (high == backend.inf or high == -backend.inf):
+                high = dtype_max if high == backend.inf else dtype_min
+            
+        if isinstance(low, (int, float)):
+            _low = backend.full([1] * len(shape), low, dtype=dtype, device=device)
+        else:
+            _low = backend.astype(low, dtype)
+        if isinstance(high, (int, float)):
+            _high = backend.full([1] * len(shape), high, dtype=dtype, device=device)
         else:
-            if isinstance(low, (int, float)):
-                _low = backend.full([1] * len(shape), low, dtype=dtype, device=device)
-            else:
-                _low = backend.astype(low, dtype)
-            if isinstance(high, (int, float)):
-                _high = backend.full([1] * len(shape), high, dtype=dtype, device=device)
-            else:
-                _high = backend.astype(high, dtype)
+            _high = backend.astype(high, dtype)
         
         _low = backend.abbreviate_array(
             _low,

unienv_interface/space/spaces/text.py

@@ -88,9 +88,7 @@ class TextSpace(Space[str, BDeviceType, BDtypeType, BRNGType]):
         return rng, sample
 
     def create_empty(self) -> str:
-        raise NotImplementedError(
-            "TextSpace does not support create_empty method. Use an empty string instead."
-        )
+        return ""
 
     def is_bounded(self, manner = "both"):
         return manner == "below" or (

unienv_interface/transformations/dict_transform.py

@@ -7,6 +7,27 @@ from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, B
 import copy
 from .transformation import DataTransformation, TargetDataT
 
+def get_chained_value(
+    data : Mapping[str, Any],
+    chained_key : List[str],
+    ignore_missing_keys : bool = False,
+) -> Any:
+    assert len(chained_key) >= 1, "Chained key must have at least one key"
+    if chained_key[0] not in data.keys():
+        if ignore_missing_keys:
+            return None
+        else:
+            raise KeyError(f"Key '{chained_key[0]}' not found in data")
+    
+    if len(chained_key) == 1:
+        return data[chained_key[0]]
+    else:
+        return get_chained_value(
+            data[chained_key[0]],
+            chained_key[1:],
+            ignore_missing_keys=ignore_missing_keys
+        )
+
 def call_function_on_chained_dict(
     data : Mapping[str, Any],
     chained_key : List[str],
@@ -111,11 +132,16 @@ class DictTransformation(DataTransformation):
     ) -> Optional["DictTransformation"]:
         if not self.has_inverse:
             return None
-        
-        inverse_mapping = {
-            key: transformation.direction_inverse(source_space)
-            for key, transformation in self.mapping.items()
-        }
+
+        inverse_mapping = {}
+        for key, transformation in self.mapping.items():
+            inverse_mapping[key] = transformation.direction_inverse(
+                None if source_space is None else get_chained_value(
+                    source_space,
+                    key.split('/'),
+                    ignore_missing_keys=self.ignore_missing_keys
+                )
+            )
         
         return DictTransformation(
             mapping=inverse_mapping,