@@ -122,35 +122,130 @@ def _assert_type_and_shape(stats_list: list[dict[str, dict]]):
122122 if "image" in fkey and k != "count" and v .shape != (3 , 1 , 1 ):
123123 raise ValueError (f"Shape of '{ k } { v .shape } )
124124
125-
126- def aggregate_feature_stats (stats_ft_list : list [dict [str , dict ]]) -> dict [str , dict [str , np .ndarray ]]:
127- """Aggregates stats for a single feature."""
128- means = np .stack ([s ["mean" ] for s in stats_ft_list ])
129- variances = np .stack ([s ["std" ] ** 2 for s in stats_ft_list ])
130- counts = np .stack ([s ["count" ] for s in stats_ft_list ])
131- total_count = counts .sum (axis = 0 )
132-
133- # Prepare weighted mean by matching number of dimensions
134- while counts .ndim < means .ndim :
135- counts = np .expand_dims (counts , axis = - 1 )
136-
137- # Compute the weighted mean
138- weighted_means = means * counts
139- total_mean = weighted_means .sum (axis = 0 ) / total_count
140-
141- # Compute the variance using the parallel algorithm
142- delta_means = means - total_mean
143- weighted_variances = (variances + delta_means ** 2 ) * counts
144- total_variance = weighted_variances .sum (axis = 0 ) / total_count
145-
146- return {
147- "min" : np .min (np .stack ([s ["min" ] for s in stats_ft_list ]), axis = 0 ),
148- "max" : np .max (np .stack ([s ["max" ] for s in stats_ft_list ]), axis = 0 ),
149- "mean" : total_mean ,
150- "std" : np .sqrt (total_variance ),
151- "count" : total_count ,
125+ def aggregate_feature_stats (stats_ft_list ):
126+ """
127+ Combine per-episode stats for a single feature.
128+
129+ Each element of stats_ft_list must be a dict with:
130+ - 'mean': array-like
131+ - 'std': array-like
132+ - 'count' or 'n': int
133+ Optional:
134+ - 'min': array-like
135+ - 'max': array-like
136+
137+ Returns:
138+ {
139+ 'mean': (D,) float64
140+ 'std': (D,) float64
141+ 'count': int
142+ # 'min': (D,) float64 (if present in any input)
143+ # 'max': (D,) float64 (if present in any input)
144+ }
145+ """
146+ def _norm (x , name , idx ):
147+ try :
148+ arr = np .asarray (x , dtype = np .float64 )
149+ except Exception as e :
150+ raise TypeError (f"{ name } { idx } { e } )
151+ arr = np .squeeze (arr )
152+ if arr .ndim == 0 :
153+ arr = arr [None ]
154+ return arr
155+
156+ means , vars_ , counts = [], [], []
157+ mins , maxs = [], []
158+ have_min = False
159+ have_max = False
160+
161+ for i , s in enumerate (stats_ft_list ):
162+ if s is None :
163+ continue
164+
165+ if "mean" not in s or ("std" not in s ):
166+ raise KeyError (f"Missing 'mean' or 'std' at idx { i } { list (s .keys ())} )
167+
168+ m = _norm (s ["mean" ], "mean" , i )
169+ sd = _norm (s ["std" ], "std" , i )
170+
171+ # count key tolerance
172+ n = s .get ("count" , s .get ("n" , None ))
173+ if n is None :
174+ raise KeyError (f"Missing 'count' (or 'n') at idx { i } )
175+ try :
176+ n = int (n )
177+ except Exception :
178+ raise TypeError (f"'count' must be int-like at idx { i } { type (n )} { n } )
179+
180+ if m .shape != sd .shape :
181+ raise ValueError (f"mean/std shape mismatch at idx { i } { m .shape } { sd .shape } )
182+
183+ means .append (m )
184+ vars_ .append (sd ** 2 )
185+ counts .append (n )
186+
187+ if "min" in s and s ["min" ] is not None :
188+ mins .append (_norm (s ["min" ], "min" , i ))
189+ have_min = True
190+ if "max" in s and s ["max" ] is not None :
191+ maxs .append (_norm (s ["max" ], "max" , i ))
192+ have_max = True
193+
194+ if not means :
195+ raise ValueError ("No valid stats provided." )
196+
197+ # Ensure all episodes have the same feature dimension
198+ shapes = {tuple (x .shape ) for x in means }
199+ if len (shapes ) != 1 :
200+ raise ValueError (f"Inconsistent feature shapes across episodes: { shapes } )
201+
202+ means = np .stack (means ) # (E, D)
203+ vars_ = np .stack (vars_ ) # (E, D)
204+ counts = np .asarray (counts , dtype = np .int64 ) # (E,)
205+
206+ N = int (counts .sum ())
207+ if N <= 1 :
208+ # Degenerate case; fall back to simple average/std
209+ pooled_mean = means .mean (axis = 0 )
210+ pooled_std = np .sqrt (vars_ .mean (axis = 0 ))
211+ else :
212+ # Weighted (pooled) mean
213+ w = counts [:, None ] # (E, 1)
214+ pooled_mean = (w * means ).sum (axis = 0 ) / N # (D,)
215+
216+ # Unbiased pooled variance:
217+ # SS_within = sum_i (n_i - 1) * var_i
218+ # SS_between = sum_i n_i * (mean_i - pooled_mean)^2
219+ ss_within = ((counts - 1 )[:, None ] * vars_ ).sum (axis = 0 )
220+ ss_between = (w * (means - pooled_mean ) ** 2 ).sum (axis = 0 )
221+ denom = max (N - 1 , 1 )
222+ pooled_var = (ss_within + ss_between ) / denom
223+ pooled_std = np .sqrt (np .maximum (pooled_var , 0.0 ))
224+
225+ out = {
226+ "mean" : pooled_mean .astype (np .float64 ),
227+ "std" : pooled_std .astype (np .float64 ),
228+ "count" : N ,
152229 }
153230
231+ # Optional min/max aggregation (elementwise)
232+ if have_min :
233+ # if some episodes lacked min/max, we only use the ones that exist
234+ mins = [m for m in mins if m is not None ]
235+ # validate shapes
236+ shapes = {tuple (np .squeeze (np .asarray (m )).shape ) for m in mins }
237+ if len (shapes ) != 1 :
238+ raise ValueError (f"Inconsistent 'min' shapes across episodes: { shapes } )
239+ out ["min" ] = np .min (np .stack (mins ), axis = 0 ).astype (np .float64 )
240+
241+ if have_max :
242+ maxs = [m for m in maxs if m is not None ]
243+ shapes = {tuple (np .squeeze (np .asarray (m )).shape ) for m in maxs }
244+ if len (shapes ) != 1 :
245+ raise ValueError (f"Inconsistent 'max' shapes across episodes: { shapes } )
246+ out ["max" ] = np .max (np .stack (maxs ), axis = 0 ).astype (np .float64 )
247+
248+ return out
154249
155250def aggregate_stats (stats_list : list [dict [str , dict ]]) -> dict [str , dict [str , np .ndarray ]]:
156251 """Aggregate stats from multiple compute_stats outputs into a single set of stats.
0 commit comments