def _reset()

in robogym/envs/rearrange/datasets/envstates/random.py [0:0]

• 35 lines of code
• 4 McCabe index (conditional complexity)

    def _reset(self, env: MeshRearrangeEnv):
        object_groups = MeshRearrangeEnv._sample_random_object_groups(env)

        num_groups = len(object_groups)
        (
            meshes,
            mesh_object_dataset_names,
            mesh_object_ids,
            mesh_scales,
        ) = self._sample_meshes(num_groups)

        attrs_per_group = {
            "material_args": MeshRearrangeEnv._sample_object_materials(env, num_groups),
            "color": MeshRearrangeEnv._sample_object_colors(env, num_groups),
            "scale": MeshRearrangeEnv._sample_object_size_scales(env, num_groups),
            "mesh_files": meshes,
        }
        attrs_per_group["scale"] *= np.array(mesh_scales)

        object_groups = MeshRearrangeEnv._set_group_attributes(
            env, object_groups, attrs_per_group
        )

        # Be sure to recreate sim before randomizing position and rotation, because
        # randomization algorithm required bounding box for the current mesh objects.
        # In addition, object_groups should be deepcopied so that scale normalization inside
        # _recreate_sim is not applied to the current variable.
        MeshRearrangeEnv._set_object_groups(env, deepcopy(object_groups))
        MeshRearrangeEnv._recreate_sim(env)

        init_quats = self._sample_object_initial_rotations(env)
        init_quats = self._post_process_quat(
            init_quats, object_groups, mesh_object_dataset_names
        )

        # Note that environment object placement function should be called after object
        # orientation is set. Therefore we first set object orientation quats to environment,
        # and then call object placement function.
        MeshRearrangeEnv._set_object_initial_rotations(env, init_quats)
        init_pos, is_valid = MeshRearrangeEnv._generate_object_placements(env)

        self.envstate = Envstate(
            object_groups=object_groups,
            object_dataset_names=mesh_object_dataset_names,
            object_ids=mesh_object_ids,
            init_pos=init_pos,
            is_valid=is_valid,
            init_quats=init_quats,
        )