def reset()

in hucc/envs/ctrlgs.py [0:0]

• 55 lines of code
• 19 McCabe index (conditional complexity)

    def reset(self):
        need_hard_reset = self._do_hard_reset or (
            self.hard_reset_interval > 0
            and self._reset_counter % self.hard_reset_interval == 0
        )
        # Reset
        if need_hard_reset:
            self.hard_reset()
            self._reset_counter = 0
        if self.relative_frame_of_reference:
            self.goal_featurizer.set_frame_of_reference()

        # Sample features and goal
        resample_features = False
        if need_hard_reset:
            resample_features = True
        if self.resample_features == 'soft':
            resample_features = True
        elif self.resample_features.startswith('soft'):
            freq = int(self.resample_features[4:])
            resample_features = self._reset_counter % freq == 0
        if resample_features:
            self._features = self.sample_features()
            self._features_s = self._feature_strings[
                ','.join(map(str, self._features))
            ]
            self._feature_mask *= 0
            for f in self._features:
                self._feature_mask[self.task_map[f]] = 1.0

        self.goal = self.sample_goals_random()[0]
        if self.goal_sampling in {'q', 'q_value'}:
            if self.model:
                self.goal = self.sample_goal_using_q()
        elif self.goal_sampling in {'r', 'reachability', 'r2', 'reachability2'}:
            if self.model:
                self.goal = self.sample_goal_using_r()
        elif self.goal_sampling not in {'random', 'uniform'}:
            raise ValueError(
                f'Unknown goal sampling method "{self.goal_sampling}"'
            )

        def distance_to_goal():
            gs = self.proj(self.goal_featurizer(), self._features)
            d = self.goal - gs
            for i, f in enumerate(self._features):
                if f in self.goal_space['twist_feats']:
                    # Wrap around projected pi/-pi for distance
                    d[i] = (
                        np.remainder(
                            (self.goal[i] - gs[i]) + self.proj_pi,
                            2 * self.proj_pi,
                        )
                        - self.proj_pi
                    )
            return np.linalg.norm(d, ord=2)

        self._d_initial = distance_to_goal()

        self._do_hard_reset = False
        self._reset_counter += 1
        self._step = 0
        return self.get_observation()