[content] input_pipeline/generate_atari_dataset.py (python)
            observations_seq.append(obs.astype(np.uint8))
            actions_seq.append(actions.astype(np.int64))

        if "final_info" in infos:
            for info in infos["final_info"]:
                if info and "episode" in info:
                    print(
                        f"global_step={global_step}, episodic_return={info['episode']['r']}"
                    )
                    writer.add_scalar(
                        "charts/episodic_return", info["episode"]["r"], global_step
                    )
                    writer.add_scalar(
                        "charts/episodic_length", info["episode"]["l"], global_step
                    )

                    continue_capturing_multi = any(
                        episodes_captured_per_split[s] < split_targets[s]
                        for s in splits_in_order
                    )
                    if args.capture_dataset and continue_capturing_multi:
                        current_len = len(observations_seq)
                        if current_len >= args.min_episode_length:
                            frames = np.concatenate(observations_seq, axis=0).astype(
                                np.uint8
                            )
                            acts = np.concatenate(actions_seq, axis=0).astype(np.int64)

                            episode_obs_chunks = []
                            episode_act_chunks = []
                            start_idx = 0
                            while start_idx < current_len:
                                end_idx = min(start_idx + args.chunk_size, current_len)
                                if end_idx - start_idx < args.chunk_size:
                                    print(
                                        f"Warning: Inconsistent chunk_sizes. Episode has {current_len} frames, "
                                        f"which is smaller than the requested chunk_size: {args.chunk_size}. "
                                        "This might lead to performance degradation during training."
                                    )
                                episode_obs_chunks.append(frames[start_idx:end_idx])
                                episode_act_chunks.append(acts[start_idx:end_idx])
                                start_idx = end_idx

                            obs_chunks_data = [
                                seq.astype(np.uint8) for seq in episode_obs_chunks
                            ]
                            act_chunks_data = [act for act in episode_act_chunks]
                            obs_chunks.extend(obs_chunks_data)
                            act_chunks.extend(act_chunks_data)

                            # Save to the active split
                            (
                                ep_metadata,
                                file_idx_by_split[current_split],
                                obs_chunks,
                                act_chunks,
                            ) = save_chunks(
                                file_idx_by_split[current_split],
                                args.chunks_per_file,
                                split_dir,
                                obs_chunks,
                                act_chunks,
                            )