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

unienv_interface/utils/control_util.py

@@ -0,0 +1,68 @@
+from fractions import Fraction
+from math import gcd
+from functools import reduce
+from typing import Iterable, Tuple, List, Union
+
+Number = Union[int, float]
+
+def _lcm(a: int, b: int) -> int:
+    return abs(a // gcd(a, b) * b) if a and b else 0
+
+def _lcmm(values: Iterable[int]) -> int:
+    return reduce(_lcm, values, 1)
+
+def _gcdm(values: Iterable[int]) -> int:
+    return reduce(gcd, values)
+
+def find_best_timestep(
+    timesteps: Iterable[Number],
+    *,
+    max_denominator: int = 10_000,
+    return_fraction: bool = False,
+) -> Tuple[Union[float, Fraction], List[int]]:
+    """
+    Compute the simulation timestep dt such that every sensor period is an
+    integer multiple of dt (i.e., dt is the GCD of the periods). Works with floats.
+
+    Args:
+        timesteps: Iterable of sensor periods (seconds, ms, etc.). Must be > 0.
+        max_denominator: Max denominator when rational-approximating floats.
+                         Increase if your periods are very fine-grained.
+        return_fraction:  If True, returns dt as a Fraction; otherwise a float.
+
+    Returns:
+        dt: The best simulation timestep (float or Fraction).
+        steps_per_sensor: For each input period T_i, the integer k_i = T_i / dt.
+
+    Raises:
+        ValueError: If list is empty or contains non-positive values.
+    """
+    # Validate and convert to Fractions
+    periods = list(timesteps)
+    if not periods:
+        raise ValueError("timesteps must be a non-empty sequence.")
+    if any(p <= 0 for p in periods):
+        raise ValueError("All timesteps must be positive.")
+
+    fracs = [Fraction(p).limit_denominator(max_denominator) for p in periods]
+
+    # Find common denominator (LCM of all denominators)
+    D = _lcmm([f.denominator for f in fracs])
+
+    # Scale each fraction to that denominator → integers
+    ints = [f.numerator * (D // f.denominator) for f in fracs]
+
+    # GCD of the integerized periods → integer g; convert back via g/D
+    g = _gcdm(ints)
+    dt_frac = Fraction(g, D)  # exact rational dt
+
+    # Sanity: dt must be > 0
+    if dt_frac <= 0:
+        raise RuntimeError("Computed non-positive timestep; check inputs.")
+
+    steps_per_sensor = [int(f // dt_frac) for f in fracs]  # exact integer division
+
+    if return_fraction:
+        return dt_frac, steps_per_sensor
+    else:
+        return float(dt_frac), steps_per_sensor

unienv_interface/utils/data_queue.py

@@ -0,0 +1,184 @@
+from typing import Dict, Any, Optional, Tuple, Union, Generic, TypeVar
+import numpy as np
+import copy
+import dataclasses
+
+from unienv_interface.space import Space
+from unienv_interface.space import batch_utils as sbu, flatten_utils as sfu
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+
+DataT = TypeVar('DataT')
+
+@dataclasses.dataclass(frozen=True)
+class SpaceDataQueueState(Generic[DataT]):
+    data: DataT # (L, B, ...) or (L, ...)
+
+    def replace(self, **changes: Any) -> 'SpaceDataQueueState':
+        return dataclasses.replace(self, **changes)
+
+class FuncSpaceDataQueue(
+    Generic[DataT, BArrayType, BDeviceType, BDtypeType, BRNGType]
+):
+    def __init__(
+        self,
+        space : Space[DataT, BDeviceType, BDtypeType, BRNGType],
+        batch_size : Optional[int],
+        maxlen: int,
+    ) -> None:
+        assert maxlen > 0, "Max length must be greater than 0"
+        assert batch_size is None or batch_size > 0, "Batch size must be greater than 0 if provided"
+        assert batch_size is None or sbu.batch_size(space) == batch_size, "Batch size must match the space's batch size if provided"
+        self.single_space = space
+        self.stacked_space = sbu.batch_space(space, maxlen) # (L, ...) or (L, B, ...)
+        self.output_space = sbu.swap_batch_dims(
+            self.stacked_space, 0, 1
+        ) if batch_size is not None else self.stacked_space # (B, L, ...) or (L, ...)
+        self._maxlen = maxlen
+        self._batch_size = batch_size
+
+    @property
+    def maxlen(self) -> int:
+        return self._maxlen
+
+    @property
+    def batch_size(self) -> Optional[int]:
+        return self._batch_size
+
+    @property
+    def backend(self) -> ComputeBackend:
+        return self.single_space.backend
+    
+    @property
+    def device(self) -> Optional[BDeviceType]:
+        return self.single_space.device
+    
+    def init(
+        self,
+        initial_data : DataT,
+    ) -> SpaceDataQueueState:
+        return self.reset(
+            SpaceDataQueueState(self.stacked_space.create_empty()),
+            initial_data
+        )
+
+    def reset(
+        self, 
+        state : SpaceDataQueueState,
+        initial_data : DataT,
+        mask : Optional[BArrayType] = None,
+    ) -> SpaceDataQueueState:
+        assert self.batch_size is None or mask is None, \
+            "Mask should not be provided if batch size is empty"
+        index = (
+            slice(None), mask
+        ) if mask is not None else slice(None)
+        
+        expanded_data = sbu.get_at( # Add a singleton horizon dimension to the data
+            self.single_space,
+            initial_data,
+            None
+        )
+        return state.replace(
+            data=sbu.set_at(
+                self.stacked_space,
+                state.data,
+                index,
+                expanded_data
+            )
+        )
+
+    def add(self, state : SpaceDataQueueState, data : DataT) -> SpaceDataQueueState:
+        new_data = self.backend.map_fn_over_arrays(
+            state.data,
+            lambda x: self.backend.roll(x, shift=-1, axis=0),
+        )
+        new_data = sbu.set_at(
+            self.stacked_space,
+            new_data,
+            -1,
+            data
+        )
+        return state.replace(data=new_data)
+    
+    def get_output_data(self, state : SpaceDataQueueState) -> DataT:
+        if self.batch_size is None:
+            return state.data
+        else:
+            return sbu.swap_batch_dims_in_data(
+                self.backend,
+                state.data,
+                0, 1
+            ) # (L, B, ...) -> (B, L, ...)
+
+class SpaceDataQueue(
+    Generic[DataT, BArrayType, BDeviceType, BDtypeType, BRNGType]
+):
+    def __init__(
+        self,
+        space : Space[DataT, BDeviceType, BDtypeType, BRNGType],
+        batch_size : Optional[int],
+        maxlen: int,
+    ) -> None:
+        self.func_queue = FuncSpaceDataQueue(
+            space,
+            batch_size,
+            maxlen
+        )
+        self.state = None
+
+    @property
+    def single_space(self) -> Space[DataT, BDeviceType, BDtypeType, BRNGType]:
+        return self.func_queue.single_space
+
+    @property
+    def stacked_space(self) -> Space[DataT, BDeviceType, BDtypeType, BRNGType]:
+        return self.func_queue.stacked_space
+
+    @property
+    def output_space(self) -> Space[DataT, BDeviceType, BDtypeType, BRNGType]:
+        return self.func_queue.output_space
+    
+    @property
+    def maxlen(self) -> int:
+        return self.func_queue.maxlen
+
+    @property
+    def batch_size(self) -> Optional[int]:
+        return self.func_queue.batch_size
+    
+
+    @property
+    def backend(self) -> ComputeBackend:
+        return self.func_queue.backend
+
+    @property
+    def device(self) -> Optional[BDeviceType]:
+        return self.func_queue.device
+
+    def reset(
+        self,
+        initial_data : DataT,
+        mask : Optional[BArrayType] = None,
+    ) -> None:
+        if self.state is None:
+            assert mask is None, "Mask should not be provided on the first reset"
+            self.state = self.func_queue.init(initial_data)
+        else:
+            self.state = self.func_queue.reset(
+                self.state,
+                initial_data,
+                mask
+            )
+        
+    def add(self, data : DataT) -> None:
+        assert self.state is not None, "Data queue must be reset before adding data"
+        self.state = self.func_queue.add(
+            self.state,
+            data
+        )
+
+    def get_output_data(self) -> DataT:
+        assert self.state is not None, "Data queue must be reset before getting output data"
+        return self.func_queue.get_output_data(
+            self.state
+        )

unienv_interface/utils/stateclass.py

@@ -0,0 +1,46 @@
+import dataclasses
+from typing import TypeVar
+import functools
+
+__all__ = ['stateclass', 'field', 'StateClass']
+
+def field(pytree_node=True, *, metadata=None, **kwargs):
+    return dataclasses.field(metadata=(metadata or {}) | {'pytree_node': pytree_node},
+                        **kwargs)
+
+def stateclass(
+    clz=None, /, **kwargs
+):
+    if clz is None:
+        return functools.partial(stateclass, **kwargs)  # type: ignore[bad-return-type]
+
+    # check if already a stateclass
+    if '_unienv_stateclass' in clz.__dict__:
+        return clz
+
+    if 'frozen' not in kwargs.keys():
+        kwargs['frozen'] = True
+    data_clz = dataclasses.dataclass(**kwargs)(clz)  # type: ignore
+    
+    def replace(self, **updates):
+        """Returns a new object replacing the specified fields with new values."""
+        return dataclasses.replace(self, **updates)
+
+    data_clz.replace = replace
+
+    # add a _unienv_stateclass flag to distinguish from regular dataclasses
+    data_clz._unienv_stateclass = True  # type: ignore[attr-defined]
+
+    return data_clz  # type: ignore
+
+TNode = TypeVar('TNode', bound='StateClass')
+
+class StateClass:
+	def __init_subclass__(cls, **kwargs):
+		stateclass(cls, **kwargs)  # pytype: disable=wrong-arg-types
+
+	def __init__(self, *args, **kwargs):
+		raise NotImplementedError
+    
+	def replace(self: TNode, **overrides) -> TNode:
+		raise NotImplementedError

unienv_interface/utils/vec_util.py

@@ -0,0 +1,15 @@
+import cloudpickle
+from typing import Callable
+
+class MultiProcessFn:
+    def __init__(self, fn : Callable):
+        self.fn = fn
+    
+    def __getstate__(self):
+        return cloudpickle.dumps(self.fn)
+    
+    def __setstate__(self, state):
+        self.fn = cloudpickle.loads(state)
+    
+    def __call__(self, *args, **kwargs):
+        return self.fn(*args, **kwargs)

unienv_interface/world/__init__.py

@@ -1,4 +1,6 @@
 from .world import World, RealWorld
 from .node import WorldNode
+from .combined_node import CombinedWorldNode
 from .funcworld import FuncWorld
-from .funcnode import FuncWorldNode
+from .funcnode import FuncWorldNode
+from .combined_funcnode import CombinedFuncWorldNode

unienv_interface/world/combined_funcnode.py

@@ -0,0 +1,336 @@
+from typing import Optional, Dict, Any, Tuple, Union, Iterable, Mapping
+
+from unienv_interface.backends import BArrayType, BDeviceType, BDtypeType, BRNGType
+from unienv_interface.space import Space, DictSpace
+
+from .funcnode import FuncWorldNode
+from .funcworld import WorldStateT
+
+CombinedDataT = Union[Dict[str, Any], Any]
+CombinedNodeStateT = Dict[str, Any]
+
+class CombinedFuncWorldNode(FuncWorldNode[
+	WorldStateT, CombinedNodeStateT,
+	Optional[CombinedDataT],  # Context type (can be None)
+	CombinedDataT,             # Observation type
+	CombinedDataT,             # Action type
+	BArrayType, BDeviceType, BDtypeType, BRNGType
+]):
+	"""A functional counterpart to `CombinedWorldNode` that composes multiple `FuncWorldNode`s.
+
+	It aggregates spaces (context, observation, action) and runtime data (context, observation, info, reward, termination, truncation)
+	across child nodes. If only one child exposes a given interface and `direct_return=True`, the value is passed through directly.
+	"""
+
+	def __init__(
+		self,
+		name: str,
+		nodes: Iterable[FuncWorldNode[WorldStateT, Any, Any, Any, Any, BArrayType, BDeviceType, BDtypeType, BRNGType]],
+		direct_return: bool = True,
+	):
+		nodes = list(nodes)
+		if len(nodes) == 0:
+			raise ValueError("At least one node is required to create a CombinedFuncWorldNode.")
+
+		first_node = nodes[0]
+		# Ensure all nodes share the same world & control timestep
+		for node in nodes[1:]:
+			assert node.world is first_node.world, "All nodes must belong to the same world." \
+				f" Mismatch between {first_node.name} and {node.name}."
+			assert node.control_timestep == first_node.control_timestep, "All nodes must have the same control timestep." \
+				f" Mismatch between {first_node.name} and {node.name}."
+
+		names = [node.name for node in nodes]
+		if len(names) != len(set(names)):
+			raise ValueError("All nodes must have unique names.")
+
+		self.nodes = nodes
+
+		# Aggregate spaces similar to `CombinedWorldNode`
+		_, self.context_space = self.aggregate_spaces(
+			{node.name: node.context_space for node in nodes if node.context_space is not None},
+			direct_return=direct_return,
+		)
+		_, self.observation_space = self.aggregate_spaces(
+			{node.name: node.observation_space for node in nodes if node.observation_space is not None},
+			direct_return=direct_return,
+		)
+		self._action_node_name_direct, self.action_space = self.aggregate_spaces(
+			{node.name: node.action_space for node in nodes if node.action_space is not None},
+			direct_return=direct_return,
+		)
+
+		self.has_reward = any(node.has_reward for node in nodes)
+		self.has_termination_signal = any(node.has_termination_signal for node in nodes)
+		self.has_truncation_signal = any(node.has_truncation_signal for node in nodes)
+
+		self.name = name
+		self.direct_return = direct_return
+
+	# ========== Helper aggregation methods ==========
+	@staticmethod
+	def aggregate_spaces(
+		spaces: Dict[str, Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]]],
+		direct_return: bool = True,
+	) -> Tuple[Optional[str], Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]]]:
+		if len(spaces) == 0:
+			return None, None
+		elif len(spaces) == 1 and direct_return:
+			return next(iter(spaces.items()))
+		else:
+			backend = next(iter(spaces.values())).backend
+			return None, DictSpace(
+				backend,
+				{name: space for name, space in spaces.items() if space is not None},
+			)
+
+	@staticmethod
+	def aggregate_data(
+		data: Dict[str, Any],
+		direct_return: bool = True,
+	) -> Optional[Union[Dict[str, Any], Any]]:
+		if len(data) == 0:
+			return None
+		elif len(data) == 1 and direct_return:
+			return next(iter(data.values()))
+		else:
+			return data
+
+	# ========== properties ==========
+	@property
+	def world(self):  # type: ignore[override]
+		return self.nodes[0].world
+
+	@property
+	def control_timestep(self):  # type: ignore[override]
+		return self.nodes[0].control_timestep
+
+	# ========== Lifecycle methods ==========
+	def initial(
+		self,
+		world_state: WorldStateT,
+		*,
+		seed: Optional[int] = None,
+		pernode_kwargs: Dict[str, Dict[str, Any]] = {},
+	) -> Tuple[WorldStateT, CombinedNodeStateT]:
+		node_states: CombinedNodeStateT = {}
+		for node in self.nodes:
+			world_state, node_state = node.initial(world_state, seed=seed, **pernode_kwargs.get(node.name, {}))
+			node_states[node.name] = node_state
+		return world_state, node_states
+
+	def reset(
+		self,
+		world_state: WorldStateT,
+		node_state: CombinedNodeStateT,
+		*,
+		seed: Optional[int] = None,
+		mask: Optional[BArrayType] = None,
+		pernode_kwargs: Dict[str, Dict[str, Any]] = {},
+		**kwargs,
+	) -> Tuple[WorldStateT, CombinedNodeStateT]:
+		node_state = node_state.copy()
+		for node in self.nodes:
+			ns = node_state[node.name]
+			world_state, ns = node.reset(
+				world_state,
+				ns,
+				seed=seed,
+				mask=mask,
+				**pernode_kwargs.get(node.name, {}),
+			)
+			node_state[node.name] = ns
+		return world_state, node_state
+
+	def after_reset(
+		self,
+		world_state: WorldStateT,
+		node_state: CombinedNodeStateT,
+		*,
+		mask: Optional[BArrayType] = None,
+	) -> Tuple[
+		WorldStateT,
+		CombinedNodeStateT,
+		Optional[CombinedDataT],
+		Optional[CombinedDataT],
+		Optional[Dict[str, Any]],
+	]:
+		node_state = node_state.copy()
+		contexts: Dict[str, Any] = {}
+		observations: Dict[str, Any] = {}
+		infos: Dict[str, Any] = {}
+
+		for node in self.nodes:
+			ns = node_state[node.name]
+			world_state, ns, ctx, obs, info = node.after_reset(world_state, ns, mask=mask)
+			node_state[node.name] = ns
+			if ctx is not None:
+				contexts[node.name] = ctx
+			if obs is not None:
+				observations[node.name] = obs
+			if info is not None:
+				infos[node.name] = info
+
+		return (
+			world_state,
+			node_state,
+			self.aggregate_data(contexts, direct_return=self.direct_return),
+			self.aggregate_data(observations, direct_return=self.direct_return),
+			self.aggregate_data(infos, direct_return=False),
+		)
+
+	def pre_environment_step(
+		self,
+		world_state: WorldStateT,
+		node_state: CombinedNodeStateT,
+		dt: Union[float, BArrayType],
+	) -> Tuple[WorldStateT, CombinedNodeStateT]:
+		node_state = node_state.copy()
+		for node in self.nodes:
+			ns = node_state[node.name]
+			world_state, ns = node.pre_environment_step(world_state, ns, dt)
+			node_state[node.name] = ns
+		return world_state, node_state
+
+	def set_next_action(
+		self,
+		world_state: WorldStateT,
+		node_state: CombinedNodeStateT,
+		action: CombinedDataT,
+	) -> Tuple[WorldStateT, CombinedNodeStateT]:
+		assert self.action_space is not None, "Action space is None, cannot set action."
+		
+		node_state = node_state.copy()
+		if self._action_node_name_direct is not None:
+			# Only one actionable node
+			for node in self.nodes:
+				if node.name == self._action_node_name_direct:
+					ns = node_state[node.name]
+					world_state, ns = node.set_next_action(world_state, ns, action)  # type: ignore[arg-type]
+					node_state[node.name] = ns
+					break
+		else:
+			assert isinstance(action, Mapping), "Action must be a mapping when there are multiple action spaces."
+			for node in self.nodes:
+				if node.action_space is not None:
+					assert node.name in action, f"Action for node {node.name} is missing."
+					ns = node_state[node.name]
+					world_state, ns = node.set_next_action(world_state, ns, action[node.name])
+					node_state[node.name] = ns
+		return world_state, node_state
+
+	def post_environment_step(
+		self,
+		world_state: WorldStateT,
+		node_state: CombinedNodeStateT,
+		dt: Union[float, BArrayType],
+	) -> Tuple[WorldStateT, CombinedNodeStateT]:
+		node_state = node_state.copy()
+		for node in self.nodes:
+			ns = node_state[node.name]
+			world_state, ns = node.post_environment_step(world_state, ns, dt)
+			node_state[node.name] = ns
+		return world_state, node_state
+
+	def close(self, world_state: WorldStateT, node_state: CombinedNodeStateT) -> WorldStateT:  # type: ignore[override]
+		for node in self.nodes:
+			world_state = node.close(world_state, node_state[node.name])
+		return world_state
+
+	# ========== Data accessors ==========
+	def get_observation(
+		self,
+		world_state: WorldStateT,
+		node_state: CombinedNodeStateT,
+	) -> CombinedDataT:
+		assert self.observation_space is not None, "Observation space is None, cannot get observation."
+		return self.aggregate_data(
+			{
+				node.name: node.get_observation(world_state, node_state[node.name])
+				for node in self.nodes
+				if node.observation_space is not None
+			},
+			direct_return=self.direct_return,
+		)
+
+	def get_reward(
+		self,
+		world_state: WorldStateT,
+		node_state: CombinedNodeStateT,
+	) -> Union[float, BArrayType]:
+		assert self.has_reward, "This node does not provide a reward."
+		if self.world.batch_size is None:
+			return sum(
+				node.get_reward(world_state, node_state[node.name])
+				for node in self.nodes
+				if node.has_reward
+			)
+		rewards = self.backend.zeros(
+			(self.world.batch_size,),
+			dtype=self.backend.default_floating_dtype,
+			device=self.device,
+		)
+		for node in self.nodes:
+			if node.has_reward:
+				rewards = rewards + node.get_reward(world_state, node_state[node.name])
+		return rewards
+
+	def get_termination(
+		self,
+		world_state: WorldStateT,
+		node_state: CombinedNodeStateT,
+	) -> Union[bool, BArrayType]:
+		assert self.has_termination_signal, "This node does not provide a termination signal."
+		if self.world.batch_size is None:
+			return any(
+				node.get_termination(world_state, node_state[node.name])
+				for node in self.nodes
+				if node.has_termination_signal
+			)
+		terminations = self.backend.zeros(
+			(self.world.batch_size,),
+			dtype=self.backend.default_boolean_dtype,
+			device=self.device,
+		)
+		for node in self.nodes:
+			if node.has_termination_signal:
+				terminations = self.backend.logical_or(
+					terminations, node.get_termination(world_state, node_state[node.name])
+				)
+		return terminations
+
+	def get_truncation(
+		self,
+		world_state: WorldStateT,
+		node_state: CombinedNodeStateT,
+	) -> Union[bool, BArrayType]:
+		assert self.has_truncation_signal, "This node does not provide a truncation signal."
+		if self.world.batch_size is None:
+			return any(
+				node.get_truncation(world_state, node_state[node.name])
+				for node in self.nodes
+				if node.has_truncation_signal
+			)
+		truncations = self.backend.zeros(
+			(self.world.batch_size,),
+			dtype=self.backend.default_boolean_dtype,
+			device=self.device,
+		)
+		for node in self.nodes:
+			if node.has_truncation_signal:
+				truncations = self.backend.logical_or(
+					truncations, node.get_truncation(world_state, node_state[node.name])
+				)
+		return truncations
+
+	def get_info(
+		self,
+		world_state: WorldStateT,
+		node_state: CombinedNodeStateT,
+	) -> Optional[Dict[str, Any]]:
+		infos: Dict[str, Any] = {}
+		for node in self.nodes:
+			info = node.get_info(world_state, node_state[node.name])
+			if info is not None:
+				infos[node.name] = info
+		return self.aggregate_data(infos, direct_return=False)  # Always dict if not empty